continuous_batching.py

# Copyright 2026 Modilify
# SPDX-License-Identifier: LicenseRef-Modilify-Open-Model-1.0
"""Continuous batching for Modilify Mk2 behind the Transformers public API shape.

The upstream continuous runner is autoregressive: it persists every query in a
paged cache and emits exactly one token per request and step.  ModilifyMk2 instead
denoises a transient bidirectional canvas and may accept a ragged token chunk.
This module consequently owns the request runner while preserving the public
manager lifecycle and ``GenerationOutput`` contract.

Accepted prefix K/V is stored per request without padding.  Every heavy denoise
step creates a temporary, left-padded batched cache view.  The decoder only
reads that view, so padding can never become persistent or evict real tokens
from a sliding-window cache.
"""

from __future__ import annotations

import asyncio
import copy
import hashlib
import math
import os
import queue
import threading
import time
import uuid
import warnings
from collections import defaultdict, deque
from collections.abc import Callable, Generator, Sequence
from dataclasses import asdict, dataclass, field, is_dataclass, replace
from typing import Any

import torch
from transformers.cache_utils import Cache, DynamicCache
from transformers.generation.configuration_utils import ContinuousBatchingConfig
from transformers.generation.continuous_batching.requests import (
    GenerationOutput,
    RequestStatus,
)

from .commit_policy import fused_commit_failure_rate, select_commit_lengths
from .generation_modilify_mk2 import (
    ModilifyMk2GenerationConfig,
    ModilifyMk2GenerationOutput,
    ModilifyMk2RollingState,
    NoiseCanvasSampler,
    _add_repetition_history,
    _flatten_token_ids,
    build_denoise_trace_event,
    deterministic_episode_iteration_bound,
)
from .latent_deliberation import (
    LatentDeliberationState,
    TrajectoryHistory,
    cat_latent_states,
    cat_trajectory_history,
    cat_trajectory_tape,
    empty_trajectory_tape,
    infer_commit_reason,
    slice_latent_state,
    slice_trajectory_history,
    slice_trajectory_tape,
)


_TERMINAL_REASONS = frozenset(
    {
        "turn_end",
        "eos",
        "max_new_tokens",
        "max_denoising_steps",
        "episode_watchdog",
        "cancelled",
        "error",
    }
)


def continuous_config_fingerprint(
    generation_config: Any,
    continuous_batching_config: ContinuousBatchingConfig | None,
) -> str:
    """Return a stable-enough in-process fingerprint for persistent reuse."""

    generation_payload = (
        generation_config.to_dict()
        if hasattr(generation_config, "to_dict")
        else vars(generation_config)
    )
    batching = continuous_batching_config or ContinuousBatchingConfig()
    batching_payload = asdict(batching) if is_dataclass(batching) else vars(batching)
    return repr(
        (
            sorted(generation_payload.items(), key=lambda item: item[0]),
            sorted(batching_payload.items(), key=lambda item: item[0]),
        )
    )


@dataclass
class ModilifyMk2ContinuousGenerationOutput(GenerationOutput):
    """Official ``GenerationOutput`` plus ModilifyMk2 request-local diagnostics."""

    stop_reason: str | None = None
    committed_tokens: int = 0
    denoise_steps: int = 0
    no_progress_steps: int = 0
    jump_count: int = 0
    forced_jump_bad_count: int = 0
    heavy_forward_count: int = 0
    latent_context_update_count: int = 0
    average_commit_len: float = 0.0
    tokens_per_forward: float = 0.0
    seed: int | None = None
    scheduler_run_id: str | None = None
    queue_seconds: float = 0.0
    inference_seconds: float = 0.0
    total_seconds: float = 0.0
    last_step_batch_size: int = 0
    is_stream_update: bool = False
    delta_tokens: list[int] = field(default_factory=list)
    state_shift_count: int = 0
    latent_memory_norm: float = 0.0
    state_retention_score: float = 0.0

    def is_finished(self) -> bool:
        """Treat failed/cancelled requests as terminal for every consumer API."""

        return self.status in {RequestStatus.FINISHED, RequestStatus.FAILED}


@dataclass
class ModilifyMk2RequestState:
    """All mutable state required to suspend and re-batch one request."""

    request_id: str
    prompt_ids: list[int]
    max_new_tokens: int
    eos_token_ids: tuple[int, ...]
    streaming: bool
    record_timestamps: bool
    seed: int
    max_denoising_steps: int | None
    trace_callback: Callable[[dict[str, object]], None] | None = None
    created_time: float = field(default_factory=time.perf_counter)
    status: RequestStatus = RequestStatus.PENDING
    started_time: float = -1.0
    finished_time: float = -1.0
    generated_tokens: list[int] = field(default_factory=list)
    logprobs: list[float] = field(default_factory=list)
    timestamps: list[float] = field(default_factory=list)
    cache: Cache | None = None
    rolling_state: ModilifyMk2RollingState | None = None
    repetition_history: torch.BoolTensor | None = None
    generator: torch.Generator | None = None
    logical_length: int = 0
    max_iterations: int = 0
    reserved_blocks: int = 0
    denoise_steps: int = 0
    jumps: int = 0
    forced_jump_tokens: int = 0
    shifts: int = 0
    stop_reason: str | None = None
    error: str | None = None
    terminal_emitted: bool = False
    last_step_batch_size: int = 0
    last_delta_tokens: list[int] = field(default_factory=list)


def _clone_tensor_row(value: torch.Tensor, row: int) -> torch.Tensor:
    return value[row : row + 1].clone()


def _slice_rolling_state(state: ModilifyMk2RollingState, row: int) -> ModilifyMk2RollingState:
    selected = slice(row, row + 1)
    return ModilifyMk2RollingState(
        canvas=_clone_tensor_row(state.canvas, row),
        confidence=_clone_tensor_row(state.confidence, row),
        entropy=_clone_tensor_row(state.entropy, row),
        age=_clone_tensor_row(state.age, row),
        latent_state=slice_latent_state(state.latent_state, selected),
        history=slice_trajectory_history(state.history, selected),
        tape=slice_trajectory_tape(state.tape, selected),
    )


def _pack_rolling_states(states: Sequence[ModilifyMk2RollingState]) -> ModilifyMk2RollingState:
    return ModilifyMk2RollingState(
        canvas=torch.cat([state.canvas for state in states], dim=0),
        confidence=torch.cat([state.confidence for state in states], dim=0),
        entropy=torch.cat([state.entropy for state in states], dim=0),
        age=torch.cat([state.age for state in states], dim=0),
        latent_state=cat_latent_states([state.latent_state for state in states]),
        history=cat_trajectory_history([state.history for state in states]),
        tape=cat_trajectory_tape([state.tape for state in states]),
    )


class ModilifyMk2LogicalCachePool:
    """Per-request hole-free cache storage with ephemeral batched read views."""

    def __init__(self, model: Any, *, max_batch_tokens: int | None = None) -> None:
        self.model = model
        self.text_config = model.config.get_text_config(decoder=True)
        self.device = model.model.decoder.embed_tokens.weight.device
        self.max_batch_tokens = max_batch_tokens

    def new_cache(self) -> DynamicCache:
        return DynamicCache(config=self.text_config)

    @torch.inference_mode()
    def prefill(self, prompt_ids: Sequence[int]) -> Cache:
        cache = self.new_cache()
        chunk_size = self.max_batch_tokens or len(prompt_ids)
        for start in range(0, len(prompt_ids), chunk_size):
            stop = min(start + chunk_size, len(prompt_ids))
            tokens = torch.tensor(
                [list(prompt_ids[start:stop])], device=self.device, dtype=torch.long
            )
            mask = torch.ones(1, stop, device=self.device, dtype=torch.bool)
            positions = torch.arange(
                start, stop, device=self.device, dtype=torch.int32
            ).unsqueeze(0)
            cache = self.model.model.encoder(
                input_ids=tokens,
                attention_mask=mask,
                past_key_values=cache,
                position_ids=positions,
            ).past_key_values
        return cache

    @torch.inference_mode()
    def append(self, state: ModilifyMk2RequestState, token_ids: Sequence[int]) -> None:
        if not token_ids:
            return
        if state.cache is None:
            raise RuntimeError("Cannot append tokens before request prefill.")
        tokens = torch.tensor([list(token_ids)], device=self.device, dtype=torch.long)
        positions = torch.arange(
            state.logical_length,
            state.logical_length + tokens.shape[1],
            device=self.device,
            dtype=torch.int32,
        ).unsqueeze(0)
        mask = torch.ones(
            1,
            state.logical_length + tokens.shape[1],
            device=self.device,
            dtype=torch.bool,
        )
        state.cache = self.model.model.encoder(
            input_ids=tokens,
            attention_mask=mask,
            past_key_values=state.cache,
            position_ids=positions,
        ).past_key_values

    def pack(
        self, states: Sequence[ModilifyMk2RequestState]
    ) -> tuple[DynamicCache, torch.BoolTensor, torch.LongTensor]:
        if not states or any(state.cache is None for state in states):
            raise ValueError("Every packed request must have an initialized cache.")
        logical_lengths = torch.tensor(
            [state.logical_length for state in states],
            device=self.device,
            dtype=torch.long,
        )
        maximum_length = int(logical_lengths.max())
        attention_mask = torch.arange(
            maximum_length, device=self.device
        )[None, :].ge(maximum_length - logical_lengths[:, None])

        packed = self.new_cache()
        source_caches = [state.cache for state in states]
        assert all(cache is not None for cache in source_caches)
        if any(len(cache.layers) != len(packed.layers) for cache in source_caches):
            raise RuntimeError("Request cache layer structures differ.")

        for layer_index, packed_layer in enumerate(packed.layers):
            source_layers = [cache.layers[layer_index] for cache in source_caches]
            if any(not layer.is_initialized for layer in source_layers):
                raise RuntimeError("Request cache contains an uninitialized layer.")
            stored_lengths = [int(layer.keys.shape[-2]) for layer in source_layers]
            maximum_stored = max(stored_lengths)

            def padded(name: str) -> torch.Tensor:
                values = []
                for layer, stored_length in zip(source_layers, stored_lengths, strict=True):
                    value = getattr(layer, name)
                    if stored_length < maximum_stored:
                        padding = value.new_zeros(
                            value.shape[0],
                            value.shape[1],
                            maximum_stored - stored_length,
                            value.shape[3],
                        )
                        value = torch.cat((padding, value), dim=-2)
                    values.append(value)
                return torch.cat(values, dim=0)

            keys = padded("keys")
            values = padded("values")
            packed_layer.lazy_initialization(keys, values)
            packed_layer.keys = keys
            packed_layer.values = values
            if hasattr(packed_layer, "cumulative_length"):
                packed_layer.cumulative_length = maximum_length
        return packed, attention_mask, logical_lengths


class ModilifyMk2ContinuousBatchingManager:
    """FIFO/prefill-first continuous manager compatible with Transformers APIs."""

    def __init__(
        self,
        model: Any,
        generation_config: ModilifyMk2GenerationConfig | None,
        continuous_batching_config: ContinuousBatchingConfig | None,
        workload_hints: Any = None,
    ) -> None:
        del workload_hints
        # Generation must not silently mutate the caller's train/eval mode.
        # Inference mode below disables autograd without changing module-local
        # dropout or other training flags.
        self.model = model
        self.generation_config = copy.deepcopy(
            generation_config or getattr(model, "generation_config", None)
            or ModilifyMk2GenerationConfig.from_model_config(model.config)
        )
        if not isinstance(self.generation_config, ModilifyMk2GenerationConfig):
            payload = self.generation_config.to_dict()
            self.generation_config = ModilifyMk2GenerationConfig(**payload)
        self.continuous_batching_config = copy.deepcopy(
            continuous_batching_config or ContinuousBatchingConfig()
        )
        self.config_fingerprint = continuous_config_fingerprint(
            self.generation_config, self.continuous_batching_config
        )
        self._validate_config()

        self.device = model.model.decoder.embed_tokens.weight.device
        self.dtype = model.model.decoder.embed_tokens.weight.dtype
        self.cache_pool = ModilifyMk2LogicalCachePool(
            model,
            max_batch_tokens=self.continuous_batching_config.max_batch_tokens,
        )
        self.sampler: NoiseCanvasSampler = model._prepare_sampler(
            self.generation_config, model.config.canvas_length
        )
        self.run_id = uuid.uuid4().hex
        self.warmed_up = False
        self.destroyed = False

        configured_requests = self.continuous_batching_config.max_requests_per_batch
        self.max_requests_per_batch = int(configured_requests or 8)
        max_batch_tokens = self.continuous_batching_config.max_batch_tokens
        if max_batch_tokens is not None:
            token_capacity = int(max_batch_tokens) // int(model.config.canvas_length)
            if token_capacity < 1:
                raise ValueError(
                    "`max_batch_tokens` must fit at least one ModilifyMk2 canvas."
                )
            self.max_requests_per_batch = min(
                self.max_requests_per_batch, token_capacity
            )
        self.block_size = int(self.continuous_batching_config.block_size)
        self.block_capacity = self._resolve_block_capacity()
        self._base_seed = (
            int(self.continuous_batching_config.seed)
            if self.continuous_batching_config.seed is not None
            else int(torch.initial_seed())
        )

        self._condition = threading.Condition(threading.RLock())
        self._pending: deque[ModilifyMk2RequestState] = deque()
        self._active: dict[str, ModilifyMk2RequestState] = {}
        self._known_request_ids: set[str] = set()
        self._cancelled: set[str] = set()
        self._output_queue: queue.Queue[ModilifyMk2ContinuousGenerationOutput] = queue.Queue()
        self._stashed_outputs: dict[
            str, deque[ModilifyMk2ContinuousGenerationOutput]
        ] = defaultdict(deque)
        self._result_handlers: dict[str, tuple[Callable, asyncio.AbstractEventLoop]] = {}
        self._thread: threading.Thread | None = None
        self._finished = threading.Event()
        self.fatal_error: BaseException | None = None
        self._input_closed = False
        self._hard_stop = False
        self._keep_for_next_session = False
        self._request_counter = 0
        self._active_reserved_blocks = 0
        self._stats = {
            "submitted": 0,
            "admitted": 0,
            "completed": 0,
            "failed": 0,
            "cancelled": 0,
            "model_steps": 0,
            "generated_tokens": 0,
            "max_observed_batch_size": 0,
            "peak_reserved_blocks": 0,
            "peak_cache_blocks": 0,
            "active_slot_steps": 0,
            "slot_capacity_steps": 0,
        }

        turn_end = self.generation_config.turn_end_token_id
        self.turn_end_token_id = int(
            model.config.turn_end_token_id if turn_end is None else turn_end
        )
        self.repetition_penalty = float(self.generation_config.repetition_penalty)
        self.excluded_repetition_token_ids = _flatten_token_ids(
            self.generation_config.repetition_penalty_exclude_token_ids,
            self.generation_config.pad_token_id,
            self.generation_config.bos_token_id,
            self.generation_config.eos_token_id,
            self.generation_config.turn_end_token_id,
            getattr(model.config, "image_token_id", None),
        )

    def _validate_config(self) -> None:
        config = self.continuous_batching_config
        positive_optional = (
            "num_blocks",
            "max_batch_tokens",
            "max_requests_per_batch",
        )
        if not isinstance(config.block_size, int) or config.block_size < 4:
            raise ValueError("`block_size` must be an integer greater than or equal to 4.")
        for name in positive_optional:
            value = getattr(config, name)
            if value is not None and (not isinstance(value, int) or value <= 0):
                raise ValueError(f"`{name}` must be a positive integer when set.")
        if config.max_blocks_per_request is not None and (
            not isinstance(config.max_blocks_per_request, int)
            or config.max_blocks_per_request < 0
        ):
            raise ValueError("`max_blocks_per_request` must be a non-negative integer.")
        if not isinstance(config.max_queue_size, int) or config.max_queue_size < 0:
            raise ValueError("`max_queue_size` must be a non-negative integer.")
        if config.scheduler_type not in {"fifo", "prefill_first"}:
            raise ValueError("ModilifyMk2 continuous batching supports `fifo` and `prefill_first`.")
        if config.max_memory_percent is not None and not (
            0.0 < float(config.max_memory_percent) <= 1.0
        ):
            raise ValueError("`max_memory_percent` must be in (0, 1].")
        if config.use_async_batching is True:
            raise ValueError("ModilifyMk2 continuous batching currently uses synchronous model steps.")
        requested_graphs = config.use_cuda_graph
        if requested_graphs is True or (
            isinstance(requested_graphs, tuple) and any(requested_graphs)
        ):
            raise ValueError("CUDA graphs are not supported by the ragged ModilifyMk2 runner.")
        if config.cpu_offload_space is not None and config.cpu_offload_space > 0:
            raise ValueError("CPU cache offload is not supported by the ModilifyMk2 runner.")
        if int(config.default_compile_level or 0) > 0:
            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
        if config.varlen_compile_config is not None or config.decode_compile_config is not None:
            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
        if config.use_default_compile_configs is True:
            raise ValueError("Continuous ModilifyMk2 compilation is not supported yet.")
        if int(config.q_padding_interval_size or 0) > 0 or int(
            config.kv_padding_interval_size or 0
        ) > 0:
            raise ValueError("Compiled continuous padding intervals are not supported.")
        if config.max_cached_graphs is not None:
            raise ValueError("Cached continuous graphs are not supported.")
        if torch.distributed.is_available() and torch.distributed.is_initialized():
            if torch.distributed.get_world_size() > 1:
                raise ValueError(
                    "Tensor/distributed parallel continuous batching is not supported."
                )
        if getattr(self.model, "device_mesh", None) is not None or getattr(
            self.model, "_device_mesh", None
        ) is not None:
            raise ValueError("Tensor-parallel continuous batching is not supported.")
        # Prefix sharing would make request ownership and row-local RNG/state
        # ambiguous.  Normalize this optimization off rather than silently use it.
        config.allow_block_sharing = False

    def _available_memory_bytes(self) -> int | None:
        if self.device.type == "cuda" and torch.cuda.is_available():
            free, _ = torch.cuda.mem_get_info(self.device)
            return int(free)
        if self.device.type == "mps" and torch.backends.mps.is_available():
            return max(
                0,
                int(torch.mps.recommended_max_memory())
                - int(torch.mps.driver_allocated_memory()),
            )
        if self.device.type == "cpu":
            try:
                import psutil

                return int(psutil.virtual_memory().available)
            except (ImportError, OSError, ValueError):
                pass
            try:
                return int(os.sysconf("SC_AVPHYS_PAGES")) * int(
                    os.sysconf("SC_PAGE_SIZE")
                )
            except (OSError, TypeError, ValueError):
                return None
        return None

    def _estimated_block_bytes(self) -> int:
        config = self.model.config.text_config
        layer_types = list(config.layer_types)
        local_heads = int(config.num_key_value_heads)
        local_dim = int(config.head_dim)
        global_heads = int(
            getattr(config, "num_global_key_value_heads", None) or local_heads
        )
        global_dim = int(getattr(config, "global_head_dim", None) or local_dim)
        per_token = 0
        for layer_type in layer_types:
            if layer_type == "full_attention":
                heads, dimension = global_heads, global_dim
            else:
                heads, dimension = local_heads, local_dim
            per_token += 2 * heads * dimension * torch.empty((), dtype=self.dtype).element_size()
        return max(1, per_token * int(self.continuous_batching_config.block_size))

    def _resolve_block_capacity(self) -> int | None:
        capacity = self.continuous_batching_config.num_blocks
        percent = self.continuous_batching_config.max_memory_percent
        available = self._available_memory_bytes()
        if percent is None and capacity is None:
            # Never make the default cache silently unbounded. This fraction is
            # applied to currently available device/host memory after model load.
            percent = 0.8
        if percent is not None and available is None:
            raise RuntimeError(
                "Cannot infer available cache memory on this device; set `num_blocks` "
                "explicitly instead of `max_memory_percent`."
            )
        if percent is not None and available is not None:
            memory_blocks = int(
                available * float(percent) / self._estimated_block_bytes()
            )
            capacity = memory_blocks if capacity is None else min(int(capacity), memory_blocks)
        return None if capacity is None else max(0, int(capacity))

    @staticmethod
    def _block_footprint(reservations: Sequence[int]) -> int:
        """Return persistent plus temporary packed-cache block equivalents."""

        if not reservations:
            return 0
        return sum(reservations) + len(reservations) * max(reservations)

    def _current_block_footprint(self) -> int:
        return self._block_footprint(
            [state.reserved_blocks for state in self._active.values()]
        )

    def _derive_seed(self, request_id: str) -> int:
        digest = hashlib.sha256(
            str(self._base_seed).encode("ascii")
            + b"\0"
            + request_id.encode("utf-8")
        ).digest()
        return int.from_bytes(digest[:8], "big") & ((1 << 63) - 1)

    @property
    def stats(self) -> dict[str, Any]:
        with self._condition:
            capacity_steps = int(self._stats["slot_capacity_steps"])
            return {
                "scheduler_run_id": self.run_id,
                **self._stats,
                "slot_utilization": (
                    float(self._stats["active_slot_steps"]) / capacity_steps
                    if capacity_steps
                    else 0.0
                ),
                "active_requests": len(self._active),
                "pending_requests": len(self._pending),
                "max_requests_per_batch": self.max_requests_per_batch,
                "block_capacity": -1 if self.block_capacity is None else self.block_capacity,
                "reserved_blocks": self._active_reserved_blocks,
                "cache_blocks": self._current_block_footprint(),
            }

    def is_running(self) -> bool:
        return self._thread is not None and self._thread.is_alive()

    def warmup(self) -> None:
        if self.destroyed:
            raise RuntimeError("Cannot warm up a destroyed manager.")
        # CUDA graphs and static-shape compilation are intentionally unsupported;
        # normal eager kernels warm naturally on the first real batch.
        self.warmed_up = True

    def start(self) -> None:
        if self._keep_for_next_session:
            self._prepare_for_next_session()
        with self._condition:
            if self.destroyed:
                raise RuntimeError("Cannot start a destroyed manager.")
            if self.is_running():
                return
            self._finished.clear()
            self.fatal_error = None
            self._hard_stop = False
            self._thread = threading.Thread(
                target=self._run_generation_loop,
                name=f"modilify_mk2-continuous-{self.run_id[:8]}",
                daemon=True,
            )
            self._thread.start()

    def join(
        self,
        stop_trigger_time: float | None = None,
        timeout: float | None = None,
    ) -> None:
        """Wait for the current worker, matching the official manager lifecycle."""

        del stop_trigger_time
        with self._condition:
            thread = self._thread
        if thread is None or thread is threading.current_thread():
            return
        thread.join(timeout=timeout)
        if thread.is_alive():
            raise TimeoutError("Timed out waiting for continuous generation to stop.")

    def _prepare_for_next_session(self) -> None:
        """Finish an asynchronous prior stop and reopen a cached manager safely."""

        with self._condition:
            if not self._keep_for_next_session:
                return
            thread = self._thread
        if thread is not None and thread.is_alive():
            thread.join()
        with self._condition:
            if self.destroyed:
                raise RuntimeError("Cannot reuse a destroyed manager.")
            if self._pending or self._active:
                raise RuntimeError("Cannot reuse a manager with unfinished requests.")
            self._input_closed = False
            self._hard_stop = False
            self._keep_for_next_session = False
            self.fatal_error = None
            self._cancelled.clear()
            self._condition.notify_all()

    def close_input(self) -> None:
        """Stop accepting requests and let the iterator drain all submitted work."""

        with self._condition:
            self._input_closed = True
            self._condition.notify_all()

    def stop(
        self,
        block: bool = True,
        timeout: float | None = None,
        keep_for_next_session: bool = False,
        hard_stop: bool = False,
    ) -> None:
        with self._condition:
            self._input_closed = True
            self._hard_stop = bool(hard_stop)
            self._keep_for_next_session = bool(keep_for_next_session)
            if hard_stop:
                self._cancelled.update(self._known_request_ids)
            self._condition.notify_all()
            thread = self._thread
        if hard_stop and (thread is None or not thread.is_alive()):
            self._apply_cancellations()
        if block and thread is not None:
            self.join(timeout=timeout)
        if keep_for_next_session and not self.is_running():
            with self._condition:
                self._input_closed = False
                self._hard_stop = False
                self._keep_for_next_session = False
                self.fatal_error = None

    def destroy(self) -> None:
        if self.destroyed:
            return
        self.stop(block=True, hard_stop=True)
        self.destroyed = True
        with self._condition:
            self._pending.clear()
            self._active.clear()
            self._condition.notify_all()

    def add_request(
        self,
        input_ids: list[int],
        request_id: str | None = None,
        max_new_tokens: int | None = None,
        streaming: bool = False,
        record_timestamps: bool = False,
        eos_token_id: int | list[int] | None = None,
        **request_kwargs: Any,
    ) -> str:
        if not input_ids or any(
            not isinstance(token_id, int) or isinstance(token_id, bool)
            for token_id in input_ids
        ):
            raise ValueError("`input_ids` must be a non-empty list of integer token IDs.")
        seed = request_kwargs.pop("seed", None)
        trace_callback = request_kwargs.pop("denoise_trace_callback", None)
        max_denoising_steps = request_kwargs.pop(
            "max_denoising_steps", self.generation_config.max_denoising_steps
        )
        if request_kwargs:
            unsupported = ", ".join(sorted(request_kwargs))
            raise ValueError(f"Unsupported per-request generation options: {unsupported}")
        if trace_callback is not None and not callable(trace_callback):
            raise TypeError("`denoise_trace_callback` must be callable.")
        if trace_callback is not None and self.max_requests_per_batch > 1:
            raise ValueError(
                "ModilifyMk2 denoise tracing remains a batch-size-1 interface; "
                "set `max_requests_per_batch=1`."
            )
        limit = self.generation_config.max_new_tokens if max_new_tokens is None else max_new_tokens
        if not isinstance(limit, int) or limit <= 0:
            raise ValueError("`max_new_tokens` must be a positive integer.")
        if max_denoising_steps is not None and (
            not isinstance(max_denoising_steps, int) or max_denoising_steps <= 0
        ):
            raise ValueError("`max_denoising_steps` must be a positive integer when set.")

        with self._condition:
            if self.destroyed or self._input_closed:
                raise RuntimeError("Continuous batching manager is not accepting requests.")
            if self.fatal_error is not None:
                raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
            if request_id is None:
                request_id = f"req_{self._request_counter}"
                self._request_counter += 1
            if request_id in self._known_request_ids:
                raise ValueError(f"Duplicate continuous request ID: {request_id}")
            queue_limit = int(self.continuous_batching_config.max_queue_size)
            deadline = time.monotonic() + 10.0
            while queue_limit and len(self._pending) >= queue_limit:
                if not self.is_running():
                    raise queue.Full(
                        "Continuous request queue is full; start the manager before "
                        "submitting more requests."
                    )
                remaining = deadline - time.monotonic()
                if remaining <= 0:
                    raise queue.Full("Continuous request queue remained full for 10 seconds.")
                self._condition.wait(timeout=remaining)
                if self.destroyed or self._input_closed:
                    raise RuntimeError(
                        "Continuous batching manager stopped while waiting for queue space."
                    )
                if self.fatal_error is not None:
                    raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
            # The worker can close/fail the manager while this producer is
            # asleep. Recheck under the same lock immediately before append.
            if self.destroyed or self._input_closed:
                raise RuntimeError("Continuous batching manager is not accepting requests.")
            if self.fatal_error is not None:
                raise RuntimeError("Continuous batching manager has failed.") from self.fatal_error
            if request_id in self._known_request_ids:
                raise ValueError(f"Duplicate continuous request ID: {request_id}")
            configured_eos = self.generation_config.eos_token_id if eos_token_id is None else eos_token_id
            if configured_eos is None:
                configured_eos = self.model.config.eos_token_id
            eos_values = (
                [configured_eos]
                if isinstance(configured_eos, int)
                else list(configured_eos or [])
            )
            stop_ids = tuple(
                dict.fromkeys(
                    [self.turn_end_token_id, *(int(value) for value in eos_values if int(value) >= 0)]
                )
            )
            resolved_seed = self._derive_seed(request_id) if seed is None else int(seed)
            state = ModilifyMk2RequestState(
                request_id=request_id,
                prompt_ids=list(input_ids),
                max_new_tokens=int(limit),
                eos_token_ids=stop_ids,
                streaming=bool(streaming),
                record_timestamps=bool(record_timestamps),
                seed=resolved_seed & ((1 << 63) - 1),
                max_denoising_steps=max_denoising_steps,
                trace_callback=trace_callback,
            )
            state.reserved_blocks = math.ceil(
                (len(state.prompt_ids) + state.max_new_tokens) / self.block_size
            )
            self._pending.append(state)
            self._known_request_ids.add(request_id)
            self._stats["submitted"] += 1
            self._condition.notify_all()
            return request_id

    def add_requests(
        self,
        inputs: list[list[int]],
        max_new_tokens: int | None = None,
        streaming: bool = False,
        record_timestamps: bool = False,
        **request_kwargs: Any,
    ) -> list[str]:
        request_ids = request_kwargs.pop("request_ids", None)
        seeds = request_kwargs.pop("seeds", None)
        if request_ids is not None and len(request_ids) != len(inputs):
            raise ValueError("`request_ids` must contain one ID per request.")
        if seeds is not None and len(seeds) != len(inputs):
            raise ValueError("`seeds` must contain one seed per request.")
        result = []
        for index, input_ids in enumerate(inputs):
            per_request = dict(request_kwargs)
            if seeds is not None:
                per_request["seed"] = seeds[index]
            result.append(
                self.add_request(
                    input_ids=input_ids,
                    request_id=None if request_ids is None else request_ids[index],
                    max_new_tokens=max_new_tokens,
                    streaming=streaming,
                    record_timestamps=record_timestamps,
                    **per_request,
                )
            )
        return result

    def cancel_request(self, request_id: str) -> None:
        with self._condition:
            if request_id in self._known_request_ids:
                self._cancelled.add(request_id)
                self._condition.notify_all()

    def register_result_handler(self, request_id: str, callback: Callable) -> None:
        loop = asyncio.get_running_loop()
        with self._condition:
            self._result_handlers[request_id] = (callback, loop)

    def _pop_stashed(self, request_id: str | None):
        with self._condition:
            if request_id is not None:
                values = self._stashed_outputs.get(request_id)
                if values:
                    return values.popleft()
                return None
            for values in self._stashed_outputs.values():
                if values:
                    return values.popleft()
        return None

    def _has_stashed_outputs(self) -> bool:
        with self._condition:
            return any(values for values in self._stashed_outputs.values())

    def get_result(
        self, request_id: str | None = None, timeout: float | None = None
    ) -> ModilifyMk2ContinuousGenerationOutput | None:
        stashed = self._pop_stashed(request_id)
        if stashed is not None:
            return stashed
        if not self.is_running() and self._output_queue.empty():
            return None
        deadline = None if timeout is None else time.monotonic() + timeout
        while True:
            remaining = None if deadline is None else max(0.0, deadline - time.monotonic())
            if remaining == 0.0:
                return None
            try:
                output = self._output_queue.get(timeout=remaining)
            except queue.Empty:
                return None
            if request_id is None or output.request_id == request_id:
                return output
            with self._condition:
                self._stashed_outputs[output.request_id].append(output)

    def __iter__(self) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]:
        while True:
            output = self.get_result(timeout=0.05)
            if output is not None:
                yield output
                continue
            if self._finished.is_set() and self._output_queue.empty():
                if not self._has_stashed_outputs():
                    return

    def request_id_iter(
        self, request_id: str
    ) -> Generator[ModilifyMk2ContinuousGenerationOutput, None, None]:
        while True:
            output = self.get_result(request_id=request_id, timeout=0.05)
            if output is not None:
                yield output
                if output.is_finished():
                    return
            elif self._finished.is_set():
                return

    def _deliver(self, output: ModilifyMk2ContinuousGenerationOutput) -> None:
        handler = None
        with self._condition:
            handler = self._result_handlers.get(output.request_id)
            if output.is_finished():
                self._result_handlers.pop(output.request_id, None)
        if handler is None:
            self._output_queue.put(output)
        else:
            callback, loop = handler
            try:
                loop.call_soon_threadsafe(callback, output)
            except RuntimeError as error:
                # A callback owner may close its event loop while a terminal
                # event is in flight. Preserve the result for pull consumers
                # instead of turning that client race into a worker fatality.
                warnings.warn(
                    f"Result callback loop closed for {output.request_id}: {error!r}",
                    stacklevel=2,
                )
                self._output_queue.put(output)

    def _output_for(
        self,
        state: ModilifyMk2RequestState,
        *,
        stream_update: bool = False,
        delta_tokens: Sequence[int] | None = None,
    ) -> ModilifyMk2ContinuousGenerationOutput:
        now = time.perf_counter()
        finished = state.status in {RequestStatus.FINISHED, RequestStatus.FAILED}
        end = state.finished_time if finished else -1.0
        shifts = max(1, state.shifts)
        steps = max(1, state.denoise_steps)
        return ModilifyMk2ContinuousGenerationOutput(
            request_id=state.request_id,
            prompt_ids=list(state.prompt_ids),
            generated_tokens=list(state.generated_tokens),
            logprobs=list(state.logprobs),
            error=state.error,
            status=state.status,
            created_time=state.created_time,
            lifespan=(state.started_time, end),
            timestamps=(list(state.timestamps) if state.record_timestamps else None),
            stop_reason=state.stop_reason,
            committed_tokens=len(state.generated_tokens),
            denoise_steps=state.denoise_steps,
            no_progress_steps=(
                0
                if state.rolling_state is None
                else int(state.rolling_state.latent_state.stagnation_steps[0])
            ),
            jump_count=state.jumps,
            forced_jump_bad_count=state.forced_jump_tokens,
            heavy_forward_count=state.denoise_steps,
            latent_context_update_count=state.denoise_steps,
            average_commit_len=len(state.generated_tokens) / shifts,
            tokens_per_forward=len(state.generated_tokens) / steps,
            seed=state.seed,
            scheduler_run_id=self.run_id,
            queue_seconds=max(0.0, state.started_time - state.created_time),
            inference_seconds=(
                max(0.0, (end if finished else now) - state.started_time)
                if state.started_time >= 0
                else 0.0
            ),
            total_seconds=max(0.0, (end if finished else now) - state.created_time),
            last_step_batch_size=state.last_step_batch_size,
            is_stream_update=stream_update,
            delta_tokens=list(
                state.last_delta_tokens if delta_tokens is None else delta_tokens
            ),
            state_shift_count=state.shifts,
            latent_memory_norm=(
                0.0
                if state.rolling_state is None
                else float(
                    state.rolling_state.latent_state.memory_slots.float()
                    .norm(dim=-1)
                    .mean()
                )
            ),
            state_retention_score=1.0 if state.shifts else 0.0,
        )

    def _finish(
        self,
        state: ModilifyMk2RequestState,
        reason: str,
        error: BaseException | str | None = None,
    ) -> None:
        if state.terminal_emitted:
            return
        if reason not in _TERMINAL_REASONS:
            raise ValueError(f"Unknown continuous stop reason: {reason}")
        state.stop_reason = reason
        state.error = None if error is None else (str(error) if isinstance(error, str) else repr(error))
        state.status = (
            RequestStatus.FAILED
            if reason in {"cancelled", "error"} or error is not None
            else RequestStatus.FINISHED
        )
        state.finished_time = time.perf_counter()
        state.terminal_emitted = True
        self._stats["completed"] += 1
        if reason == "cancelled":
            self._stats["cancelled"] += 1
        elif error is not None:
            self._stats["failed"] += 1
        self._deliver(self._output_for(state))

    def _fail_all_requests(self, error: BaseException) -> None:
        """Convert an unexpected worker failure into one terminal result per request."""

        self.fatal_error = error
        with self._condition:
            pending = list(self._pending)
            active = list(self._active.values())
            self._pending.clear()
            self._active.clear()
            self._active_reserved_blocks = 0
            self._input_closed = True
            for state in [*active, *pending]:
                self._finish(state, "error", error)
            self._condition.notify_all()

    def _request_fits(self, state: ModilifyMk2RequestState) -> bool:
        per_request_limit = self.continuous_batching_config.max_blocks_per_request
        if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit:
            return False
        if self.block_capacity is None:
            return True
        reservations = [
            *(active.reserved_blocks for active in self._active.values()),
            state.reserved_blocks,
        ]
        return self._block_footprint(reservations) <= self.block_capacity

    def _request_can_ever_fit(self, state: ModilifyMk2RequestState) -> bool:
        per_request_limit = self.continuous_batching_config.max_blocks_per_request
        if per_request_limit not in (None, 0) and state.reserved_blocks > per_request_limit:
            return False
        return (
            self.block_capacity is None
            or self._block_footprint([state.reserved_blocks]) <= self.block_capacity
        )

    def _initialize_request(self, state: ModilifyMk2RequestState) -> None:
        generator = torch.Generator(device=self.device)
        generator.manual_seed(state.seed)
        state.generator = generator
        try:
            canvas = self.sampler.initialize_canvas(
                1, self.device, generators=[generator]
            )
        except TypeError:
            canvas = self.sampler.initialize_canvas(1, self.device)
        dtype = self.model.model.decoder.embed_tokens.weight.dtype
        canvas_length = int(self.model.config.canvas_length)
        latent = LatentDeliberationState.empty(
            batch_size=1,
            canvas_length=canvas_length,
            latent_dim=self.model.config.latent_dim,
            memory_slots=self.model.config.latent_memory_slots,
            device=self.device,
            dtype=dtype,
        )
        state.rolling_state = ModilifyMk2RollingState(
            canvas=canvas,
            confidence=torch.zeros(1, canvas_length, device=self.device, dtype=torch.float32),
            entropy=torch.full(
                (1, canvas_length),
                math.log(self.model.config.text_config.vocab_size),
                device=self.device,
                dtype=torch.float32,
            ),
            age=torch.zeros(1, canvas_length, device=self.device, dtype=torch.int32),
            latent_state=latent,
            history=TrajectoryHistory.empty(
                batch_size=1,
                canvas_length=canvas_length,
                hidden_size=self.model.config.text_config.hidden_size,
                history_length=self.model.config.latent_history_length,
                device=self.device,
                dtype=dtype,
            ),
            tape=empty_trajectory_tape(
                batch_size=1,
                config=self.model.config,
                device=self.device,
                dtype=dtype,
            ),
        )
        state.cache = self.cache_pool.prefill(state.prompt_ids)
        state.logical_length = len(state.prompt_ids)
        if self.repetition_penalty != 1.0:
            state.repetition_history = torch.zeros(
                self.model.config.text_config.vocab_size,
                device=self.device,
                dtype=torch.bool,
            )
            prompt = torch.tensor([state.prompt_ids], device=self.device, dtype=torch.long)
            _add_repetition_history(
                state.repetition_history.unsqueeze(0),
                prompt,
                torch.ones_like(prompt, dtype=torch.bool),
                self.excluded_repetition_token_ids,
            )
        state.max_iterations = deterministic_episode_iteration_bound(
            torch.tensor([state.max_new_tokens]),
            max_ponder_steps=self.generation_config.max_ponder_steps,
        )
        state.started_time = time.perf_counter()
        state.status = RequestStatus.DECODING

    def _apply_cancellations(self) -> None:
        with self._condition:
            cancelled = set(self._cancelled)
            self._cancelled.clear()
            if not cancelled:
                return
            retained = deque()
            while self._pending:
                state = self._pending.popleft()
                if state.request_id in cancelled:
                    self._finish(state, "cancelled", "request cancelled")
                else:
                    retained.append(state)
            self._pending = retained
            for request_id in cancelled:
                state = self._active.pop(request_id, None)
                if state is not None:
                    self._active_reserved_blocks -= state.reserved_blocks
                    self._finish(state, "cancelled", "request cancelled")
            self._condition.notify_all()

    def _admit_requests(self) -> None:
        while True:
            with self._condition:
                if len(self._active) >= self.max_requests_per_batch or not self._pending:
                    return
                state = self._pending[0]
                if not self._request_can_ever_fit(state):
                    self._pending.popleft()
                    self._finish(
                        state,
                        "error",
                        "request exceeds continuous cache block limits",
                    )
                    self._condition.notify_all()
                    continue
                if not self._request_fits(state):
                    return
                self._pending.popleft()
                self._condition.notify_all()
            try:
                self._initialize_request(state)
            except Exception as error:
                self._finish(state, "error", error)
                continue
            with self._condition:
                if state.request_id in self._cancelled:
                    self._cancelled.remove(state.request_id)
                    self._finish(state, "cancelled", "request cancelled")
                    continue
                self._active[state.request_id] = state
                self._active_reserved_blocks += state.reserved_blocks
                self._stats["admitted"] += 1
                self._stats["peak_reserved_blocks"] = max(
                    self._stats["peak_reserved_blocks"],
                    self._active_reserved_blocks,
                )
                self._stats["peak_cache_blocks"] = max(
                    self._stats["peak_cache_blocks"],
                    self._current_block_footprint(),
                )
                self._condition.notify_all()

    def _select_rowwise_policy(
        self,
        states: Sequence[ModilifyMk2RequestState],
        proposal: torch.LongTensor,
        normal_failure_rate: torch.Tensor,
        previous_failure_rate: torch.Tensor,
        greedy_proposal: torch.LongTensor,
        jump_failure_rate: torch.Tensor,
        rolling: ModilifyMk2RollingState,
    ):
        decisions = []
        for row, state in enumerate(states):
            remaining = state.max_new_tokens - len(state.generated_tokens)
            decisions.append(
                select_commit_lengths(
                    sampled_token_ids=proposal[row : row + 1],
                    normal_failure_rate=normal_failure_rate[row : row + 1],
                    previous_failure_rate=previous_failure_rate[row : row + 1],
                    greedy_token_ids=greedy_proposal[row : row + 1],
                    jump_failure_rate=jump_failure_rate[row : row + 1],
                    ponder_steps=rolling.latent_state.ponder_steps[row : row + 1],
                    stagnation_steps=rolling.latent_state.stagnation_steps[row : row + 1],
                    active_rows=torch.ones(1, device=self.device, dtype=torch.bool),
                    remaining_lengths=torch.tensor([remaining], device=self.device),
                    failure_budget=self.generation_config.commit_failure_budget,
                    jump_failure_budget=self.generation_config.jump_failure_budget,
                    stop_token_id=state.eos_token_ids,
                    max_ponder_steps=self.generation_config.max_ponder_steps,
                    stagnation_threshold=self.generation_config.jump_on_no_progress_after,
                    min_progress=self.generation_config.min_trajectory_progress,
                )
            )
        return (
            torch.cat([decision.normal_lengths for decision in decisions]),
            torch.cat([decision.commit_lengths for decision in decisions]),
            torch.cat([decision.commit_token_ids for decision in decisions]),
            torch.cat([decision.jump_rows for decision in decisions]),
            torch.cat([decision.ponder_steps for decision in decisions]),
            torch.cat([decision.stagnation_steps for decision in decisions]),
        )

    @torch.inference_mode()
    def _run_batch_step(self, states: Sequence[ModilifyMk2RequestState]) -> list[str]:
        started = time.perf_counter()
        rolling_states = [state.rolling_state for state in states]
        if any(state is None for state in rolling_states):
            raise RuntimeError("Active request has no rolling state.")
        rolling = _pack_rolling_states(rolling_states)  # type: ignore[arg-type]
        packed_cache, cache_mask, logical_lengths = self.cache_pool.pack(states)
        batch_size = len(states)
        canvas_length = int(self.model.config.canvas_length)
        decoder_positions = (
            logical_lengths[:, None]
            + torch.arange(canvas_length, device=self.device)[None, :]
        ).to(torch.int32)
        decoder_mask = torch.cat(
            (
                cache_mask,
                torch.ones(
                    batch_size,
                    canvas_length,
                    device=self.device,
                    dtype=torch.bool,
                ),
            ),
            dim=-1,
        )
        repetition_history = None
        if self.repetition_penalty != 1.0:
            repetition_history = torch.stack(
                [state.repetition_history for state in states], dim=0  # type: ignore[list-item]
            )
        generators = [state.generator for state in states]
        if any(generator is None for generator in generators):
            raise RuntimeError("Active request has no sampling generator.")

        output = self.model(
            input_ids=None,
            past_key_values=packed_cache,
            decoder_input_ids=rolling.canvas,
            previous_confidence=rolling.confidence,
            previous_entropy=rolling.entropy,
            token_age=rolling.age,
            latent_state=rolling.latent_state,
            history=rolling.history,
            tape=rolling.tape,
            decoder_position_ids=decoder_positions,
            decoder_read_cache=True,
            decoder_attention_mask=decoder_mask,
            compact_vocab=True,
            denoise_temperature=self.generation_config.denoise_temperature,
            repetition_token_mask=repetition_history,
            repetition_penalty=self.repetition_penalty,
            sampling_generators=generators,
        )
        required = (
            output.proposal,
            output.proposal_confidence,
            output.token_entropy,
            output.greedy_proposal,
            output.greedy_confidence,
            output.next_latent_state,
        )
        if any(value is None for value in required):
            raise RuntimeError("Compact ModilifyMk2 forward did not return proposal state.")
        proposal = output.proposal
        proposal_confidence = output.proposal_confidence
        token_entropy = output.token_entropy
        greedy_proposal = output.greedy_proposal
        greedy_confidence = output.greedy_confidence
        next_canvas = proposal.clone()
        next_confidence = proposal_confidence.float()
        next_latent = replace(
            output.next_latent_state,
            confidence=next_confidence.detach().float(),
            entropy=token_entropy.detach().float(),
            age=rolling.age + 1,
            token_changed=next_canvas.ne(rolling.canvas).detach().float(),
            confidence_delta=next_confidence.detach().float() - rolling.confidence,
            entropy_delta=token_entropy.detach().float() - rolling.entropy,
        )
        live_mask = torch.ones(
            rolling.canvas.shape, device=rolling.canvas.device, dtype=torch.bool
        )
        tape_probes, tape_valid = self.model.latent_deliberation.encode_tape_frame(
            output.heavy_hidden_state, live_mask
        )
        next_state = ModilifyMk2RollingState(
            canvas=next_canvas,
            confidence=next_confidence,
            entropy=token_entropy,
            age=rolling.age + 1,
            latent_state=next_latent,
            history=rolling.history.append(
                output.heavy_hidden_state,
                next_confidence,
                token_entropy,
                next_canvas.ne(rolling.canvas).detach().float(),
                live_mask=live_mask,
            ),
            tape=rolling.tape.append(tape_probes, tape_valid),
        )
        normal_failure_rate = fused_commit_failure_rate(
            proposal_confidence,
            token_entropy,
            vocab_size=self.model.config.text_config.vocab_size,
        )
        jump_failure_rate = fused_commit_failure_rate(
            greedy_confidence,
            token_entropy,
            vocab_size=self.model.config.text_config.vocab_size,
        )
        previous_failure_rate = fused_commit_failure_rate(
            rolling.confidence,
            rolling.entropy,
            vocab_size=self.model.config.text_config.vocab_size,
        )
        (
            normal_commit,
            commit_lengths,
            commit_token_ids,
            jump_rows,
            next_ponder,
            next_stagnation,
        ) = self._select_rowwise_policy(
            states,
            proposal,
            normal_failure_rate,
            previous_failure_rate,
            greedy_proposal,
            jump_failure_rate,
            rolling,
        )
        positions = torch.arange(canvas_length, device=self.device)[None, :]
        commit_positions = positions.lt(commit_lengths[:, None])
        policy_prefix_mask = positions.lt(normal_commit[:, None])
        if bool(jump_rows.any()):
            next_state = replace(
                next_state,
                canvas=torch.where(
                    commit_positions & jump_rows[:, None],
                    commit_token_ids,
                    next_state.canvas,
                ),
            )
        next_state = replace(
            next_state,
            latent_state=replace(
                next_state.latent_state,
                ponder_steps=next_ponder,
                stagnation_steps=next_stagnation,
            ),
        )
        unshifted_trace_states = [
            _slice_rolling_state(next_state, row) for row in range(batch_size)
        ]
        if output.history_projected is None or output.working_state is None:
            raise RuntimeError("Forward did not return working trajectory features.")
        next_state = self.model._write_committed_memory(
            previous_history=rolling.history,
            next_state=next_state,
            working_state=output.working_state,
            history_projected=output.history_projected,
            heavy_hidden=output.heavy_hidden_state,
            commit_lengths=commit_lengths,
            prefix_lengths=logical_lengths,
            commit_reason=infer_commit_reason(
                commit_lengths,
                jump_rows=jump_rows,
                commit_token_ids=commit_token_ids,
                terminal_token_ids=getattr(
                    self.model.config, "terminal_token_ids", ()
                ),
            ),
        )
        shifted = self.model._shift_state_rows(
            next_state,
            commit_lengths,
            self.sampler,
            generators=generators,
        )
        shifted_states = [
            _slice_rolling_state(shifted, row) for row in range(batch_size)
        ]
        selected_confidence = torch.where(
            jump_rows[:, None], greedy_confidence, proposal_confidence
        ).float()

        finished_ids = []
        for row, state in enumerate(states):
            state.last_step_batch_size = batch_size
            state.denoise_steps += 1
            commit_length = int(commit_lengths[row])
            chunk = commit_token_ids[row, :commit_length].detach().cpu().tolist()
            state.last_delta_tokens = [int(token_id) for token_id in chunk]
            before = len(state.generated_tokens)
            try:
                self.cache_pool.append(state, chunk)
            except Exception as error:
                self._finish(state, "error", error)
                finished_ids.append(state.request_id)
                continue
            state.logical_length += commit_length
            state.generated_tokens.extend(int(token_id) for token_id in chunk)
            if self.continuous_batching_config.return_logprobs and commit_length:
                probabilities = selected_confidence[row, :commit_length].clamp_min(
                    torch.finfo(torch.float32).tiny
                )
                state.logprobs.extend(probabilities.log().detach().cpu().tolist())
            if state.record_timestamps and commit_length:
                state.timestamps.extend([time.perf_counter()] * commit_length)
            if state.repetition_history is not None and commit_length:
                tokens = commit_token_ids[row : row + 1]
                eligible = commit_positions[row : row + 1]
                _add_repetition_history(
                    state.repetition_history.unsqueeze(0),
                    tokens,
                    eligible,
                    self.excluded_repetition_token_ids,
                )
            state.jumps += int(jump_rows[row])
            if bool(jump_rows[row]):
                state.forced_jump_tokens += commit_length
            if commit_length:
                state.shifts += 1
            state.rolling_state = shifted_states[row]

            reason = None
            if self.turn_end_token_id in chunk:
                reason = "turn_end"
            elif any(token_id in state.eos_token_ids for token_id in chunk):
                reason = "eos"
            elif len(state.generated_tokens) >= state.max_new_tokens:
                reason = "max_new_tokens"
            elif (
                state.max_denoising_steps is not None
                and state.denoise_steps >= state.max_denoising_steps
            ):
                reason = "max_denoising_steps"
            elif state.denoise_steps >= state.max_iterations:
                reason = "episode_watchdog"

            elapsed = time.perf_counter() - started
            if state.trace_callback is not None:
                trace = build_denoise_trace_event(
                    denoise_step=state.denoise_steps,
                    prefix_length=state.logical_length - commit_length,
                    committed_before=before,
                    committed_after=len(state.generated_tokens),
                    no_progress_steps=int(next_stagnation[row]),
                    policy_prefix_mask=policy_prefix_mask[row : row + 1],
                    commit_length=commit_length,
                    ponder_fallback=bool(jump_rows[row]),
                    state=unshifted_trace_states[row],
                    proposal=proposal[row : row + 1],
                    committed_token_ids=commit_token_ids[row : row + 1, :commit_length],
                    step_elapsed_seconds=elapsed,
                    latent_residual_diagnostics=None,
                )
                trace["request_id"] = state.request_id
                trace["batch_size"] = batch_size
                try:
                    state.trace_callback(trace)
                except Exception as error:
                    warnings.warn(
                        f"Denoise trace callback failed for {state.request_id}: {error!r}",
                        stacklevel=2,
                    )

            if reason is not None:
                self._finish(state, reason)
                finished_ids.append(state.request_id)
            elif state.streaming and commit_length:
                self._deliver(
                    self._output_for(
                        state,
                        stream_update=True,
                        delta_tokens=state.last_delta_tokens,
                    )
                )

        self._stats["model_steps"] += 1
        self._stats["generated_tokens"] += int(commit_lengths.sum())
        self._stats["max_observed_batch_size"] = max(
            self._stats["max_observed_batch_size"], batch_size
        )
        self._stats["active_slot_steps"] += batch_size
        self._stats["slot_capacity_steps"] += self.max_requests_per_batch
        return finished_ids

    def _run_step_with_isolation(self, states: Sequence[ModilifyMk2RequestState]) -> None:
        generator_states = {
            state.request_id: state.generator.get_state()
            for state in states
            if state.generator is not None
        }
        try:
            finished_ids = self._run_batch_step(states)
        except Exception as batch_error:
            for state in states:
                if state.generator is not None:
                    state.generator.set_state(generator_states[state.request_id])
            if len(states) == 1:
                self._finish(states[0], "error", batch_error)
                finished_ids = [states[0].request_id]
            else:
                finished_ids = []
                for state in states:
                    if state.terminal_emitted:
                        finished_ids.append(state.request_id)
                        continue
                    try:
                        finished_ids.extend(self._run_batch_step([state]))
                    except Exception as request_error:
                        self._finish(state, "error", request_error)
                        finished_ids.append(state.request_id)
        with self._condition:
            for request_id in dict.fromkeys(finished_ids):
                state = self._active.pop(request_id, None)
                if state is not None:
                    self._active_reserved_blocks -= state.reserved_blocks
            self._condition.notify_all()

    @torch.inference_mode()
    def _run_generation_loop(self) -> None:
        try:
            while True:
                self._apply_cancellations()
                if self._hard_stop:
                    self._apply_cancellations()
                with self._condition:
                    has_active = bool(self._active)
                # ``prefill_first`` fills every available slot before the next
                # denoise step.  FIFO lets the already-active cohort take its
                # next step first, then fills slots released by that step.
                if (
                    self.continuous_batching_config.scheduler_type == "prefill_first"
                    or not has_active
                ):
                    self._admit_requests()
                with self._condition:
                    active = list(self._active.values())
                    should_finish = (
                        self._input_closed and not self._pending and not active
                    )
                    if should_finish:
                        return
                    if not active:
                        self._condition.wait(timeout=0.05)
                        continue
                self._run_step_with_isolation(active)
                if self.continuous_batching_config.scheduler_type == "fifo":
                    self._admit_requests()
        except BaseException as error:
            self._fail_all_requests(error)
        finally:
            self._finished.set()
            with self._condition:
                self._condition.notify_all()


@torch.inference_mode()
def generate_static_batch_with_logical_cache(
    model: Any,
    input_ids: torch.LongTensor,
    attention_mask: torch.BoolTensor | None,
    generation_config: ModilifyMk2GenerationConfig,
    *,
    seeds: Sequence[int] | None = None,
    max_new_tokens: Sequence[int] | None = None,
) -> ModilifyMk2GenerationOutput:
    """Run one fixed cohort through the same hole-free continuous engine."""

    batch_size, input_width = input_ids.shape
    if attention_mask is None:
        attention_mask = torch.ones_like(input_ids, dtype=torch.bool)
    else:
        attention_mask = attention_mask.to(device=input_ids.device, dtype=torch.bool)
    if attention_mask.shape != input_ids.shape:
        raise ValueError("`attention_mask` must have the same shape as `input_ids`.")
    prompts = [
        input_ids[row, attention_mask[row]].detach().cpu().tolist()
        for row in range(batch_size)
    ]
    if any(not prompt for prompt in prompts):
        raise ValueError("Every batched ModilifyMk2 prompt must contain at least one token.")
    if seeds is not None and len(seeds) != batch_size:
        raise ValueError("`seeds` must contain one seed per batch row.")
    if max_new_tokens is None:
        max_new_tokens = [int(generation_config.max_new_tokens)] * batch_size
    if len(max_new_tokens) != batch_size or any(
        not isinstance(limit, int) or isinstance(limit, bool) or limit <= 0
        for limit in max_new_tokens
    ):
        raise ValueError("`max_new_tokens` must contain one positive limit per batch row.")

    batching_config = ContinuousBatchingConfig(
        block_size=max(4, int(getattr(model.config, "kv_cache_bucket_size", 128))),
        max_batch_tokens=batch_size * int(model.config.canvas_length),
        max_requests_per_batch=batch_size,
        allow_block_sharing=False,
        scheduler_type="prefill_first",
    )
    manager = ModilifyMk2ContinuousBatchingManager(
        model=model,
        generation_config=generation_config,
        continuous_batching_config=batching_config,
    )
    try:
        request_ids = []
        for row, prompt in enumerate(prompts):
            request_kwargs = {}
            if seeds is not None:
                request_kwargs["seed"] = int(seeds[row])
            request_ids.append(
                manager.add_request(
                    prompt,
                    request_id=f"static_{row}",
                    max_new_tokens=int(max_new_tokens[row]),
                    streaming=False,
                    max_denoising_steps=generation_config.max_denoising_steps,
                    eos_token_id=generation_config.eos_token_id,
                    **request_kwargs,
                )
            )
        manager.close_input()
        # A static batch is one fixed cohort: queue every row before the worker
        # starts so its first heavy forward necessarily contains the full batch.
        manager.start()
        final = {}
        for output in manager:
            if output.is_finished():
                final[output.request_id] = output
        ordered = [final[request_id] for request_id in request_ids]
    finally:
        manager.stop(block=True, hard_stop=True)
        manager.destroy()
    failures = [output for output in ordered if output.error is not None]
    if failures:
        details = "; ".join(
            f"{output.request_id}: {output.error}" for output in failures
        )
        raise RuntimeError(f"Static ModilifyMk2 batch generation failed: {details}")

    lengths = torch.tensor(
        [len(output.generated_tokens) for output in ordered],
        device=input_ids.device,
        dtype=torch.long,
    )
    output_width = int(lengths.max()) if lengths.numel() else 0
    pad_token_id = generation_config.pad_token_id
    if isinstance(pad_token_id, (list, tuple)):
        pad_token_id = pad_token_id[0]
    pad_token_id = int(0 if pad_token_id is None else pad_token_id)
    generated = torch.full(
        (batch_size, output_width),
        pad_token_id,
        device=input_ids.device,
        dtype=input_ids.dtype,
    )
    for row, output in enumerate(ordered):
        if output.generated_tokens:
            generated[row, : len(output.generated_tokens)] = torch.tensor(
                output.generated_tokens,
                device=input_ids.device,
                dtype=input_ids.dtype,
            )

    def tensor(name: str, *, dtype: torch.dtype) -> torch.Tensor:
        return torch.tensor(
            [getattr(output, name) for output in ordered],
            device=input_ids.device,
            dtype=dtype,
        )

    return ModilifyMk2GenerationOutput(
        sequences=torch.cat((input_ids, generated), dim=-1),
        generated_lengths=lengths,
        tokens_per_forward=tensor("tokens_per_forward", dtype=torch.float32),
        past_key_values=None,
        stop_reason=tuple(output.stop_reason for output in ordered),
        committed_tokens=lengths.clone(),
        denoise_steps=tensor("denoise_steps", dtype=torch.long),
        no_progress_steps=tensor("no_progress_steps", dtype=torch.long),
        jump_count=tensor("jump_count", dtype=torch.long),
        forced_jump_bad_count=tensor("forced_jump_bad_count", dtype=torch.long),
        heavy_forward_count=tensor("heavy_forward_count", dtype=torch.long),
        latent_context_update_count=tensor(
            "latent_context_update_count", dtype=torch.long
        ),
        average_commit_len=tensor("average_commit_len", dtype=torch.float32),
        state_shift_count=tensor("state_shift_count", dtype=torch.long),
        latent_memory_norm=tensor("latent_memory_norm", dtype=torch.float32),
        state_retention_score=tensor("state_retention_score", dtype=torch.float32),
    )


__all__ = [
    "ModilifyMk2ContinuousBatchingManager",
    "ModilifyMk2ContinuousGenerationOutput",
    "ModilifyMk2LogicalCachePool",
    "ModilifyMk2RequestState",
    "continuous_config_fingerprint",
    "generate_static_batch_with_logical_cache",
]