Upload folder using huggingface_hub

serving/backend/README.md

@@ -0,0 +1,29 @@
+# HybridAgent vLLM backend (T3)
+
+vLLM plugin that serves `togethercomputer/HybridAgent-8B-Agentic`
+(Qwen3-8B-Base backbone + HybridLink latent adapter) inside vLLM 0.30 on the
+forked core (`/opt/hybridagent-vllm/vllm`, branch `hybridagent-latent-fork`).
+
+- `backend/model.py` — `HybridAgentForCausalLM`: wraps the in-tree
+  `Qwen3ForCausalLM` and adds the `HybridLink` MLP. Provides the
+  `project_latent(hidden)` hook the fork's `gpu_model_runner` calls to feed
+  `R(last_hidden)` as `inputs_embeds` on zero-token latent steps. Backbone
+  weights load through the stock Qwen3 path, so an EXL3 checkpoint
+  (`quant_method: "exl3"`, OrcaSAQ2 linear layers) loads as a stock Qwen3
+  would. `adapter.pt` loads with `torch.load(weights_only=True)` + strict
+  `load_state_dict`.
+- `backend/hybridlink.py` — verbatim port of `HybridLink` from
+  togethercomputer/HybridAgent (parameter names match `adapter.pt`).
+- `backend/__init__.py` — `register()` entry point
+  (`vllm.general_plugins`): maps the checkpoint arch string
+  (`Qwen3ForCausalLM`, per the released config.json — re-confirm at T7) to
+  `HybridAgentForCausalLM` and arms the OrcaSAQ2 `exl3` quant config when
+  installed.
+
+Install into the serving venv to activate:
+
+    pip install -e /opt/hybridagent-vllm/backend
+
+Synthetic validation (no model download):
+
+    /opt/hybridagent-vllm/venv/bin/python /opt/hybridagent-vllm/poc/test_backend_synthetic.py

serving/backend/backend/__init__.py

@@ -0,0 +1,57 @@
+"""vLLM plugin: HybridAgent backend (Qwen3 backbone + HybridLink latent MLP).
+
+Wired through the `vllm.general_plugins` entry point (see pyproject.toml):
+
+    [project.entry-points."vllm.general_plugins"]
+    hybridagent_backend = "backend:register"
+
+vLLM calls `register()` in every process that resolves model classes
+(engine front-end and the registry's inspector subprocesses), so the
+architecture used by HybridAgent checkpoints maps to the latent-capable
+backend everywhere.
+
+The arch string is the one in the released model's config.json
+("architectures": ["Qwen3ForCausalLM"], verified against
+togethercomputer/HybridAgent-8B-Agentic). Re-confirm it against the exact
+checkpoint at T7. Registering under it intentionally shadows the in-tree
+Qwen3 model in this isolated venv: the venv serves HybridAgent only.
+"""
+from .model import HybridAgentForCausalLM
+
+ARCH = "Qwen3ForCausalLM"
+
+
+def register() -> None:
+    """Entry point: register the HybridAgent backend model architecture."""
+    from vllm.logger import init_logger
+    from vllm.model_executor.models.registry import ModelRegistry
+
+    logger = init_logger("vllm.hybridagent")
+    ModelRegistry.register_model(ARCH, "backend.model:HybridAgentForCausalLM")
+    logger.info(
+        "hybridagent backend armed: %s -> HybridAgentForCausalLM "
+        "(zero-token latent steps via the project_latent hook)",
+        ARCH,
+    )
+    _arm_exl3(logger)
+
+
+def _arm_exl3(logger) -> None:
+    """Reuse the OrcaSAQ2 'exl3' quantization config when it is installed.
+
+    The backbone's EXL3 linear layers come from the existing OrcaSAQ2
+    plugin (config.json `quant_method: "exl3"`); this backend does not
+    reinvent them. Without it, only dense checkpoints can load here.
+    """
+    try:
+        import orcasaq2
+    except ImportError:
+        logger.warning(
+            "hybridagent backend: orcasaq2 not installed in this environment; "
+            "EXL3 checkpoints cannot load here (dense checkpoints are unaffected)"
+        )
+        return
+    orcasaq2.register()
+
+
+__all__ = ["ARCH", "HybridAgentForCausalLM", "register"]

serving/backend/backend/hybridlink.py

@@ -0,0 +1,29 @@
+"""HybridLink: a lightweight residual adapter that maps a last-layer hidden
+state back into the input-embedding space,
+R(h) = post_ln(h + proj2(gelu(proj1(pre_ln(h)))))
+
+Ported verbatim from togethercomputer/HybridAgent (`hybrid_agent/hybridlink.py`
+@ HEAD 1b41cebb). The parameter names match the released checkpoints
+(`adapter.pt`), so the file loads with a strict `load_state_dict`.
+"""
+from __future__ import annotations
+
+import torch
+import torch.nn as nn
+
+# Tag stored in `adapter_config.json` next to the HybridLink weights.
+HYBRIDLINK_CONFIG = {"adapter_type": "ln_res_adapter"}
+
+
+class HybridLink(nn.Module):
+    def __init__(self, hidden_size: int) -> None:
+        super().__init__()
+        self.proj1 = nn.Linear(hidden_size, hidden_size)
+        self.act = nn.GELU()
+        self.proj2 = nn.Linear(hidden_size, hidden_size)
+        self.pre_ln = nn.LayerNorm(hidden_size)
+        self.post_ln = nn.LayerNorm(hidden_size)
+
+    def forward(self, h: torch.Tensor) -> torch.Tensor:
+        out = self.proj2(self.act(self.proj1(self.pre_ln(h))))
+        return self.post_ln(h + out)

serving/backend/backend/model.py

@@ -0,0 +1,96 @@
+"""vLLM backend for HybridAgent: Qwen3 backbone + HybridLink latent adapter.
+
+Wraps vLLM's in-tree Qwen3ForCausalLM and adds the HybridLink MLP as a
+submodule. The forked gpu_model_runner (T2) detects latent-capable models
+via the `project_latent` hook and feeds R(last_hidden) as inputs_embeds for
+zero-token latent steps; this class provides that hook and keeps the stock
+Qwen3 forward/logits/weight-loading contract everywhere else.
+
+Weight loading:
+  * Backbone: the stock Qwen3 path (AutoWeightsLoader over the checkpoint's
+    safetensors). Because the in-tree layers are built from
+    `vllm_config.quant_config`, an EXL3 checkpoint (config.json
+    `quant_method: "exl3"`, served by the OrcaSAQ2 linear layers) loads
+    exactly as a stock Qwen3 would -- this backend does not reinvent EXL3.
+  * HybridLink: the separate `adapter.pt` file next to the checkpoint,
+    loaded with `torch.load(..., weights_only=True)` + strict
+    `load_state_dict`, mirroring HybridAgentModel.from_pretrained.
+"""
+from __future__ import annotations
+
+import json
+from pathlib import Path
+
+import torch
+
+from vllm.config import VllmConfig
+from vllm.logger import init_logger
+from vllm.model_executor.models.qwen3 import Qwen3ForCausalLM
+
+from .hybridlink import HYBRIDLINK_CONFIG, HybridLink
+
+logger = init_logger(__name__)
+
+HYBRIDLINK_WEIGHTS = "adapter.pt"
+HYBRIDLINK_CONFIG_FILE = "adapter_config.json"
+
+
+class HybridAgentForCausalLM(Qwen3ForCausalLM):
+    """Qwen3ForCausalLM + HybridLink, for the forked zero-token latent runner."""
+
+    def __init__(self, *, vllm_config: VllmConfig, prefix: str = "") -> None:
+        super().__init__(vllm_config=vllm_config, prefix=prefix)
+        hidden_size = self.config.hidden_size
+        # S4: the HybridLink MLP stays BF16 regardless of backbone dtype.
+        self.hybridlink = HybridLink(hidden_size).to(torch.bfloat16)
+        self._compute_dtype = vllm_config.model_config.dtype
+
+    def project_latent(self, hidden: torch.Tensor) -> torch.Tensor:
+        """R(h): the hook the forked gpu_model_runner calls for latent steps.
+
+        Accepts a [D] or [B, D] last-layer hidden state and returns the same
+        shape in the model's compute dtype, ready to be consumed as
+        inputs_embeds for the zero-token latent position.
+        """
+        x = hidden
+        link_dtype = next(self.hybridlink.parameters()).dtype
+        if x.dtype != link_dtype:
+            x = x.to(link_dtype)
+        # no_grad: the hook is called inside execute_model's inference_mode,
+        # but this keeps it correct from any context (LayerNorm on inference
+        # tensors fails under an autograd-tracked mode).
+        with torch.no_grad():
+            out = self.hybridlink(x)
+        if out.dtype != self._compute_dtype:
+            out = out.to(self._compute_dtype)
+        return out
+
+    def load_weights(self, weights):
+        loaded = super().load_weights(weights)
+        self._load_hybridlink()
+        # The MLP weights come from adapter.pt, not the checkpoint stream;
+        # report them as loaded so the loader's completeness tracker passes.
+        loaded.update(f"hybridlink.{name}" for name, _ in self.hybridlink.named_parameters())
+        return loaded
+
+    def _load_hybridlink(self) -> None:
+        """Strict-load the latent adapter from adapter.pt next to the checkpoint."""
+        model_dir = Path(self.vllm_config.model_config.model)
+        weights_path = model_dir / HYBRIDLINK_WEIGHTS
+        if not weights_path.is_file():
+            raise FileNotFoundError(
+                f"{HYBRIDLINK_WEIGHTS} not found in {model_dir}; "
+                "HybridAgent serving requires the latent adapter"
+            )
+        config_path = model_dir / HYBRIDLINK_CONFIG_FILE
+        if config_path.is_file():
+            tag = json.loads(config_path.read_text(encoding="utf-8")).get(
+                "adapter_type"
+            )
+            if tag != HYBRIDLINK_CONFIG["adapter_type"]:
+                raise ValueError(
+                    f"unsupported HybridLink config {tag!r} in {config_path}"
+                )
+        state = torch.load(str(weights_path), map_location="cpu", weights_only=True)
+        self.hybridlink.load_state_dict(state, strict=True)
+        logger.info("hybridagent backend: HybridLink loaded from %s", weights_path)

serving/backend/pyproject.toml

@@ -0,0 +1,18 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hybridagent-vllm-backend"
+version = "0.1.0"
+description = "vLLM backend for HybridAgent: Qwen3ForCausalLM backbone + HybridLink latent MLP (zero-token latent steps)"
+requires-python = ">=3.10"
+# vllm and torch are intentionally not declared: they come from the serving
+# environment, mirroring orcasaq2-kernel.
+dependencies = []
+
+[project.entry-points."vllm.general_plugins"]
+hybridagent_backend = "backend:register"
+
+[tool.setuptools]
+packages = ["backend"]

serving/fork/v1/core/kv_cache_manager.py

@@ -0,0 +1,960 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import itertools
+from collections.abc import Sequence
+from dataclasses import dataclass
+from typing import Literal, overload
+
+from vllm.distributed.kv_events import BlockStored, KVCacheEvent
+from vllm.logger import init_logger
+from vllm.utils.math_utils import cdiv
+from vllm.v1.core.kv_cache_coordinator import (
+    HybridKVCacheCoordinator,
+    get_kv_cache_coordinator,
+)
+from vllm.v1.core.kv_cache_metrics import KVCacheMetricsCollector
+from vllm.v1.core.kv_cache_utils import KVCacheBlock, KVCacheBlockCopy
+from vllm.v1.core.single_type_kv_cache_manager import MambaManager
+from vllm.v1.kv_cache_interface import (
+    AttentionSpec,
+    CrossAttentionSpec,
+    EncoderOnlyAttentionSpec,
+    KVCacheConfig,
+    MambaSpec,
+    get_kv_cache_spec_kind,
+    get_kv_cache_spec_sliding_window,
+)
+from vllm.v1.metrics.stats import PrefixCacheStats
+from vllm.v1.request import Request, RequestStatus
+
+logger = init_logger(__name__)
+
+
+@dataclass
+class KVCacheBlocks:
+    """
+    The allocation result of KVCacheManager, work as the interface between
+    Scheduler and KVCacheManager, to hide KVCacheManager's internal data
+    structure from the Scheduler.
+    """
+
+    blocks: tuple[Sequence[KVCacheBlock], ...]
+    """
+    `blocks[i][j]` refers to the i-th kv_cache_group
+    and the j-th block of tokens.We don't use block of
+    tokens as the outer dimension because it assumes all
+    kv_cache_groups have the same number of blocks, which is true for now but
+    will be broken if we want to give different block_size to different
+    kv_cache_groups in the future.
+
+    Each single type KVCacheBlocks could be represented as:
+    - list[KVCacheBlock] for more than one KVCacheBlock
+    - an empty tuple for requests without KVCacheBlock
+      (a precomputed KVCacheBlocks is in KVCacheManager to avoid GC overhead)
+    """
+
+    def __add__(self, other: "KVCacheBlocks") -> "KVCacheBlocks":
+        """Adds two KVCacheBlocks instances."""
+        return KVCacheBlocks(
+            tuple(
+                list(itertools.chain(blk1, blk2))
+                for blk1, blk2 in zip(self.blocks, other.blocks)
+            )
+        )
+
+    @overload
+    def get_block_ids(
+        self,
+        allow_none: Literal[False] = False,
+        *,
+        group_ids: Sequence[int] | None = None,
+    ) -> tuple[list[int], ...]: ...
+
+    @overload
+    def get_block_ids(
+        self,
+        allow_none: Literal[True] = True,
+        *,
+        group_ids: Sequence[int] | None = None,
+    ) -> tuple[list[int], ...] | None: ...
+
+    def get_block_ids(
+        self,
+        allow_none: bool = False,
+        *,
+        group_ids: Sequence[int] | None = None,
+    ) -> tuple[list[int], ...] | None:
+        """
+        Converts the KVCacheBlocks instance to block_ids.
+
+        Args:
+            allow_none: Return None when every selected group is empty.
+            group_ids: KV cache groups to include. Includes all groups by default.
+
+        Returns:
+            tuple[list[int], ...]: A tuple of lists where:
+                - the outer tuple corresponds to KV cache groups
+                - each inner list contains the block_ids of the blocks in that
+                  group
+        """
+        groups = (
+            self.blocks
+            if group_ids is None
+            else tuple(self.blocks[group_id] for group_id in group_ids)
+        )
+        if allow_none and all(len(group) == 0 for group in groups):
+            return None
+        return tuple([blk.block_id for blk in group] for group in groups)
+
+    def get_unhashed_block_ids(self) -> list[int]:
+        """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
+        assert len(self.blocks) == 1, "Only one group is supported"
+        return [block.block_id for block in self.blocks[0] if block.block_hash is None]
+
+    def get_unhashed_block_ids_all_groups(self) -> list[list[int]]:
+        """Get block_ids of unhashed blocks from KVCacheBlocks instance."""
+        # Skip padding blocks.
+        return [
+            [
+                block.block_id
+                for block in group
+                if block.block_hash is None and not block.is_null
+            ]
+            for group in self.blocks
+        ]
+
+    def new_empty(self) -> "KVCacheBlocks":
+        """
+        Creates a new KVCacheBlocks instance with no blocks.
+        """
+        return KVCacheBlocks(tuple(() for _ in range(len(self.blocks))))
+
+
+class KVCacheManager:
+    def __init__(
+        self,
+        kv_cache_config: KVCacheConfig,
+        max_model_len: int,
+        scheduler_block_size: int,
+        hash_block_size: int,
+        max_in_flight_tokens: int | None = None,
+        enable_caching: bool = True,
+        use_eagle: bool = False,
+        num_prefill_lookahead: int = 0,
+        log_stats: bool = False,
+        enable_kv_cache_events: bool = False,
+        dcp_world_size: int = 1,
+        pcp_world_size: int = 1,
+        metrics_collector: KVCacheMetricsCollector | None = None,
+        watermark: float = 0.0,
+        enable_mamba_fine_grained_prefix_cache: bool = False,
+    ) -> None:
+        self.max_model_len = max_model_len
+        # When unset, fall back to `max_model_len` so the recycling-aware cap
+        # collapses to the prior (uncapped) admission behavior. The scheduler
+        # always supplies the real value at runtime.
+        if max_in_flight_tokens is None:
+            max_in_flight_tokens = max_model_len
+
+        self.enable_caching = enable_caching
+        self.enable_kv_cache_events = enable_kv_cache_events
+        self.use_eagle = use_eagle
+        self.log_stats = log_stats
+        self.metrics_collector = metrics_collector
+        # FIXME: make prefix cache stats conditional on log_stats. We still need
+        # this comment because when the log stats is enabled there are still
+        # potential configs we could expose in the future.
+        self.prefix_cache_stats = PrefixCacheStats() if log_stats else None
+
+        self.coordinator = get_kv_cache_coordinator(
+            kv_cache_config=kv_cache_config,
+            max_model_len=self.max_model_len,
+            max_in_flight_tokens=max_in_flight_tokens,
+            use_eagle=self.use_eagle,
+            enable_caching=self.enable_caching,
+            enable_kv_cache_events=enable_kv_cache_events,
+            dcp_world_size=dcp_world_size,
+            pcp_world_size=pcp_world_size,
+            scheduler_block_size=scheduler_block_size,
+            hash_block_size=hash_block_size,
+            metrics_collector=self.metrics_collector,
+            num_prefill_lookahead=num_prefill_lookahead,
+        )
+        # One predicate, read by both sides of the feature, so the scheduler
+        # cannot end a chunk at a junction the manager would refuse -- a refused
+        # junction costs a forward pass and displaces the block-boundary stop.
+        # Multi-module MTP is excluded because ``cache_blocks`` then hands the
+        # manager ``num_computed - num_reprefillable`` rather than the chunk end.
+        self.mamba_fine_grained_prefix_cache = (
+            enable_mamba_fine_grained_prefix_cache
+            and bool(self.coordinator.eagle_group_ids)
+            and self.coordinator.enable_partial_hash_hits
+            and self.coordinator.num_reprefillable_tokens == 0
+        )
+        if self.mamba_fine_grained_prefix_cache:
+            for manager in self.coordinator.single_type_managers:
+                if isinstance(manager, MambaManager):
+                    manager.fine_grained_prefix_cache = True
+        self.num_kv_cache_groups = len(kv_cache_config.kv_cache_groups)
+        self.block_pool = self.coordinator.block_pool
+        self.retained_hit_group_ids = tuple(
+            manager.kv_cache_group_id
+            for manager in self.coordinator.single_type_managers
+            if manager.retains_longer_hit
+        )
+        self.kv_cache_config = kv_cache_config
+
+        # Watermark: minimum number of KV cache blocks to keep free when
+        # admitting waiting/preempted requests, to avoid frequent preemptions.
+        assert watermark >= 0.0, "watermark must be non-negative"
+        self.watermark_blocks = int(watermark * kv_cache_config.num_blocks)
+        self.kv_cache_event_metadata = tuple(
+            (
+                get_kv_cache_spec_kind(group.kv_cache_spec).value,
+                get_kv_cache_spec_sliding_window(group.kv_cache_spec),
+            )
+            for group in kv_cache_config.kv_cache_groups
+        )
+
+        # Pre-constructed KVCacheBlocks with no blocks, callers should use this
+        # via create_kv_cache_blocks instead of creating new ones to avoid GC
+        # overhead.
+        #
+        # We use nested tuples to ensure the empty KVCacheBlocks is immutable.
+        self.empty_kv_cache_blocks = KVCacheBlocks(
+            tuple(() for _ in range(self.num_kv_cache_groups))
+        )
+
+    @property
+    def usage(self) -> float:
+        """Get the KV cache usage.
+
+        Returns:
+            The KV cache usage (between 0.0 and 1.0).
+        """
+        return self.block_pool.get_usage()
+
+    def make_prefix_cache_stats(self) -> PrefixCacheStats | None:
+        """Get (and reset) the prefix cache stats.
+
+        Returns:
+            The current prefix caching stats, or None if logging is disabled.
+        """
+        if not self.log_stats:
+            return None
+        stats = self.prefix_cache_stats
+        self.prefix_cache_stats = PrefixCacheStats()
+        return stats
+
+    def prefix_cache_lookup_enabled(self, request: Request) -> bool:
+        """Whether a local prefix cache lookup may be run for this request."""
+        return self.enable_caching and not request.skip_reading_prefix_cache
+
+    def record_prefix_cache_stats(self, request: Request, num_hits: int) -> None:
+        # Don't count a request that skipped the cache lookup.
+        if not self.log_stats or not self.prefix_cache_lookup_enabled(request):
+            return
+        assert self.prefix_cache_stats is not None
+        self.prefix_cache_stats.record(
+            num_tokens=request.num_tokens,
+            num_hits=num_hits,
+            preempted=request.num_preemptions > 0,
+        )
+
+    def get_computed_blocks(self, request: Request) -> tuple[KVCacheBlocks, int, int]:
+        """Get the computed (cached) blocks for the request.
+        Note that the computed blocks must be full.
+
+        Args:
+            request: The request to get the computed blocks.
+
+        Returns:
+            A tuple containing:
+                - A list of blocks that are computed for the request.
+                - The number of computed tokens.
+                - ``shared_prefix_boundary``: the block-aligned token position of
+                  a shared prefix that a sparse-retention group (Mamba / sliding
+                  window) has not cached yet (Marconi-style APC), or 0 if none.
+                  Pinned so sparse prefix-cache retention does not drop
+                  the junction and defeat cross-request reuse.
+        """
+        # We skip finding the prefix cache hit when prefix caching is
+        # disabled or the request is marked as skipping kv cache read
+        # (which happens when the request requires prompt logprobs
+        # or calls a pooling model with all pooling).
+        if not self.prefix_cache_lookup_enabled(request):
+            return self.empty_kv_cache_blocks, 0, 0
+
+        # NOTE: When all tokens hit the cache, we must recompute the last token
+        # to obtain logits. Thus, set max_cache_hit_length to prompt_length - 1.
+        # This can trigger recomputation of an entire block, rather than just
+        # the single last token, because allocate_slots() requires
+        # num_computed_tokens to be block-size aligned. Removing this limitation
+        # could slightly improve performance in the future.
+        max_cache_hit_length = request.num_tokens - 1
+        computed_blocks, num_new_computed_tokens, num_uncached = (
+            self.coordinator.find_longest_cache_hit(
+                request.block_hashes, max_cache_hit_length
+            )
+        )
+
+        # When kv_cache_report_mode is "full", emit BlockStored events
+        # for the reused prefix cache blocks so that external consumers
+        # (e.g. gateway) can learn about them.
+        if (
+            num_new_computed_tokens > 0
+            and self.enable_kv_cache_events
+            and getattr(request, "kv_cache_report_mode", "incremental") == "full"
+        ):
+            self.coordinator.emit_cached_block_events(request, computed_blocks)
+
+        # The junction to pin is where the lagging sparse-retention group stops
+        # (``num_new_computed_tokens``) plus the uncached shared prefix -- i.e.
+        # the longest single-group hit. Sub-block gaps are left to the mask,
+        # which floors to the alignment boundary (a no-op there).
+        shared_prefix_boundary = (
+            num_new_computed_tokens + num_uncached if num_uncached else 0
+        )
+
+        blocks = self.create_kv_cache_blocks(computed_blocks)
+        return blocks, num_new_computed_tokens, shared_prefix_boundary
+
+    def get_computed_blocks_for_connector(
+        self, request: Request
+    ) -> tuple[KVCacheBlocks, int, int, bool]:
+        """Local prefix-cache lookup for a request scheduled with a KV connector.
+
+        Hybrid (Mamba + full-attention) models can have per-group prefix hits
+        diverge under block pressure: the full-attention tail may be evicted
+        while a deeper Mamba state survives, or vice versa. Report the
+        full-attention hit as the local prefix - the connector transfers the
+        remaining suffix and the Mamba state is transferred unconditionally by
+        nixl's ``_apply_prefix_caching`` - and flag when that hit ran deeper
+        than a lagging group. Such a hit only has a valid Mamba state at its
+        boundary if the connector supplies it, so the caller must fall back to
+        ``get_computed_blocks`` to reconcile when no external tokens are found.
+
+        Non-hybrid models and already-convergent hits use ``get_computed_blocks``.
+
+        Returns:
+            The ``get_computed_blocks`` triple (blocks, number of local computed
+            tokens, shared-prefix boundary) plus ``hit_diverged``.
+        """
+        coordinator = self.coordinator
+        if not (
+            self.kv_cache_config.has_mamba_layers
+            and isinstance(coordinator, HybridKVCacheCoordinator)
+            and coordinator.full_attention_group_id is not None
+        ):
+            return *self.get_computed_blocks(request), False
+
+        if not self.prefix_cache_lookup_enabled(request):
+            return self.empty_kv_cache_blocks, 0, 0, False
+
+        fa_group_id = coordinator.full_attention_group_id
+        computed, per_group_hits = coordinator.find_longest_cache_hit_per_group(
+            request.block_hashes, request.num_tokens - 1
+        )
+        if any(hit > per_group_hits[fa_group_id] for hit in per_group_hits):
+            # A lagging group hit deeper than full attention means its
+            # full-attention blocks were evicted; use the reconciled boundary
+            # that every group agrees on.
+            return *self.get_computed_blocks(request), False
+
+        num_local = per_group_hits[fa_group_id]
+        blocks = self.create_kv_cache_blocks(computed)
+        # Per-group lookups do not detect an uncached shared prefix (boundary 0).
+        return blocks, num_local, 0, min(per_group_hits) < num_local
+
+    def allocate_slots(
+        self,
+        request: Request,
+        num_new_tokens: int,
+        num_new_computed_tokens: int = 0,
+        new_computed_blocks: KVCacheBlocks | None = None,
+        num_lookahead_tokens: int = 0,
+        num_external_computed_tokens: int = 0,
+        delay_cache_blocks: bool = False,
+        num_encoder_tokens: int = 0,
+        full_sequence_must_fit: bool = False,
+        reserved_blocks: int = 0,
+        has_scheduled_reqs: bool = True,
+        num_latent_positions: int = 0,  # FORK: zero-token latent KV positions
+    ) -> KVCacheBlocks | None:
+        """Add slots for a request with new tokens to append.
+
+        Args:
+            request: The request to allocate slots.
+            num_new_tokens: The number of new tokens to be allocated and computed.
+            num_new_computed_tokens: The number of new computed tokens just
+                hitting the prefix caching, excluding external tokens.
+            new_computed_blocks: The cached blocks for the above new computed
+                tokens, grouped as a tuple by kv cache groups.
+            num_lookahead_tokens: The number of speculative tokens to allocate.
+                This is used by spec decode proposers with kv-cache such
+                as eagle.
+            num_external_computed_tokens: The number of tokens that their
+                KV caches are not cached by vLLM but cached by the connector.
+            delay_cache_blocks: Whether to skip caching the blocks. This is
+                used by P/D when allocating blocks used in a KV transfer
+                which will complete in a future step.
+            num_encoder_tokens: The number of encoder tokens to allocate for
+                cross-attention in encoder-decoder models(e.g., Whisper).
+                For decoder-only models, this should be 0.
+            full_sequence_must_fit: Only allocate blocks if the KV cache has enough
+                free blocks to hold the full sequence, accounting for prefix cache hits
+                and sliding window. Used as an admission gate to prevent over-admitting
+                requests when chunked prefill would otherwise only check the first chunk
+            reserved_blocks: Number of free blocks that must be left available for
+                other in-flight sequences to complete. The actual allocation is only
+                made if it fits within (free blocks - reserved_blocks). Used to gate
+                async KV-connector loads so their initial allocation cannot consume
+                blocks an already in-flight (prefilling) sequence is relying on.
+            has_scheduled_reqs: Whether any requests are already scheduled to run
+                this step, controls whether watermark is applied.
+            num_latent_positions: FORK: number of zero-token latent KV positions
+                to allocate for this step (HybridAgent latent steps consume a KV
+                position without a token id).
+
+        Blocks layout:
+        ```
+        ----------------------------------------------------------------------
+        | < comp > | < new_comp > | < ext_comp >  | < new >  | < lookahead > |
+        ----------------------------------------------------------------------
+                                                  |   < to be computed >     |
+        ----------------------------------------------------------------------
+                                  |            < to be allocated >           |
+        ----------------------------------------------------------------------
+                                  | < to be cached (roughly, |
+                                  | details below)>          |
+        ----------------------------------------------------------------------
+        | Prefix-cached tokens from either vLLM   |
+        | or connector. Can be safely removed if  |
+        | they are outside sliding window.        |
+        ----------------------------------------------------------------------
+        |   < cached by vLLM >    | not cached by |
+                                  | vLLM, but     |
+        | ref_cnt  | ref_cnt not  | cached by     |
+        | increased| increased yet| connector     |
+        ----------------------------------------------------------------------
+        ```
+
+        Abbrivations:
+
+        ```
+        comp      = request.num_computed_tokens
+        new_comp  = num_new_computed_tokens
+                  = len(new_computed_blocks) * block_size
+        ext_comp  = num_external_computed_tokens, cached by the connector
+        new       = num_new_tokens, including unverified draft tokens
+        lookahead = num_lookahead_tokens
+        ```
+
+        NOTE: for new tokens which include both verified and unverified draft
+        tokens, we only cache the verified tokens (by capping the number at
+        `request.num_tokens`).
+
+        The allocation has three stages:
+        - Free unnecessary blocks in `comp` and check
+           if we have sufficient free blocks (return None if not).
+        - Handle prefix tokens (`comp + new_comp + ext_comp`):
+            - Free unnecessary blocks (e.g. outside sliding window)
+            - Allocate new blocks for `ext_comp` tokens inside
+              sliding window
+        - Allocate new blocks for tokens to be computed (`new + lookahead`)
+
+        Returns:
+            A list of new allocated blocks.
+        """
+        # When loading KV data asynchronously, we may have zero new tokens to
+        # compute while still allocating slots for externally computed tokens.
+        # FORK: a zero-token latent step (num_latent_positions > 0) also
+        # allocates slots without any new token id.
+        if (
+            num_new_tokens == 0
+            and num_external_computed_tokens == 0
+            and num_latent_positions == 0
+        ):
+            raise ValueError(
+                "num_new_tokens must be greater than 0 when there are no "
+                "external computed tokens"
+            )
+
+        if new_computed_blocks is not None:
+            new_computed_block_list = new_computed_blocks.blocks
+            if self.retained_hit_group_ids:
+                # Only retain the prefix covered by the local and external hits.
+                reused = num_new_computed_tokens + num_external_computed_tokens
+                groups = list(new_computed_block_list)
+                for group_id in self.retained_hit_group_ids:
+                    manager = self.coordinator.single_type_managers[group_id]
+                    groups[group_id] = groups[group_id][
+                        : cdiv(reused, manager.block_size)
+                    ]
+                new_computed_block_list = tuple(groups)
+        else:
+            new_computed_block_list = self.empty_kv_cache_blocks.blocks
+
+        # The number of computed tokens is the number of computed tokens plus
+        # the new prefix caching hits
+        num_local_computed_tokens = (
+            request.num_computed_tokens + num_new_computed_tokens
+        )
+        total_computed_tokens = min(
+            num_local_computed_tokens + num_external_computed_tokens,
+            self.max_model_len,
+        )
+
+        watermark_blocks = 0
+        # The watermark is applied to waiting/preempted requests only, and only
+        # when there's at least one request already scheduled.
+        if has_scheduled_reqs and request.status in (
+            RequestStatus.WAITING,
+            RequestStatus.PREEMPTED,
+        ):
+            watermark_blocks = self.watermark_blocks
+
+        if full_sequence_must_fit:
+            # First check and fail if the full request sequence won't fit.
+            full_num_tokens = min(request.num_tokens, self.max_model_len)
+
+            num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
+                request_id=request.request_id,
+                num_tokens=full_num_tokens,
+                new_computed_blocks=new_computed_block_list,
+                num_encoder_tokens=num_encoder_tokens,
+                total_computed_tokens=total_computed_tokens,
+                num_local_computed_tokens=num_local_computed_tokens,
+                num_tokens_main_model=full_num_tokens,
+                apply_admission_cap=True,
+            )
+            required_blocks = num_blocks_to_allocate + watermark_blocks
+            if required_blocks > self.block_pool.get_num_free_blocks():
+                return None
+
+        # FORK: latent positions are KV positions too, so the block table must
+        # extend to cover them even though no token id exists for them.
+        num_tokens_main_model = (
+            total_computed_tokens + num_new_tokens + num_latent_positions
+        )
+        num_tokens_need_slot = min(
+            num_tokens_main_model + num_lookahead_tokens, self.max_model_len
+        )
+
+        # Free the blocks that are skipped during the attention computation
+        # (e.g., tokens outside the sliding window).
+        # We can do this even if we cannot schedule this request due to
+        # insufficient free blocks.
+        # Should call this function before allocating new blocks to reduce
+        # the number of evicted blocks.
+        # Free on the processed-token basis: in-flight steps' attention windows
+        # still read blocks below the optimistic boundary, and rejected spec
+        # tokens can roll it back.
+        self.coordinator.remove_skipped_blocks(
+            request.request_id,
+            max(0, total_computed_tokens - request.num_in_flight_tokens),
+            num_prompt_tokens=request.num_prompt_tokens,
+        )
+
+        num_blocks_to_allocate = self.coordinator.get_num_blocks_to_allocate(
+            request_id=request.request_id,
+            num_tokens=num_tokens_need_slot,
+            new_computed_blocks=new_computed_block_list,
+            num_encoder_tokens=num_encoder_tokens,
+            total_computed_tokens=num_local_computed_tokens
+            + num_external_computed_tokens,
+            num_local_computed_tokens=num_local_computed_tokens,
+            num_tokens_main_model=num_tokens_main_model,
+        )
+
+        # Keep `reserved_blocks` free for other in-flight sequences, and an
+        # additional watermark of headroom for waiting/preempted admissions.
+        available_blocks = self.block_pool.get_num_free_blocks() - reserved_blocks
+        required_blocks = num_blocks_to_allocate + watermark_blocks
+        if required_blocks > available_blocks:
+            # Cannot allocate new blocks
+            return None
+
+        if (
+            new_computed_block_list is not self.empty_kv_cache_blocks.blocks
+            or num_external_computed_tokens > 0
+        ):
+            # Append the new computed blocks to the request blocks until now to
+            # avoid the case where the new blocks cannot be allocated.
+            self.coordinator.allocate_new_computed_blocks(
+                request_id=request.request_id,
+                new_computed_blocks=new_computed_block_list,
+                num_local_computed_tokens=num_local_computed_tokens,
+                num_external_computed_tokens=num_external_computed_tokens,
+            )
+
+        new_blocks = self.coordinator.allocate_new_blocks(
+            request.request_id,
+            num_tokens_need_slot,
+            num_tokens_main_model,
+            num_encoder_tokens,
+        )
+
+        # P/D: delay caching blocks if we have to recv from
+        # remote. Update state for locally cached blocks.
+        if delay_cache_blocks:
+            return self.create_kv_cache_blocks(new_blocks)
+
+        # NOTE(woosuk): We want to commit (cache) up to num_local_computed_tokens
+        # + num_external_computed_tokens + num_new_tokens, but must exclude
+        # "non-committable" tokens (e.g., draft tokens that could be rejected).
+        # Therefore, we cap the number at `request.num_tokens`, ensuring only
+        # "finalized" tokens are cached.
+        # FORK: never cache beyond the frozen cacheable prefix. A latent
+        # block's KV derives from an injected embedding, not a token id, so it
+        # must never enter the prefix cache (and BlockPool.cache_full_blocks
+        # would index past request.block_hashes if it did).
+        num_tokens_to_cache = min(
+            total_computed_tokens + num_new_tokens,
+            request.num_tokens,
+            request.kv_cacheable_prefix or self.max_model_len,
+        )
+        self.coordinator.cache_blocks(request, num_tokens_to_cache)
+
+        return self.create_kv_cache_blocks(new_blocks)
+
+    def free(self, request: Request) -> None:
+        """Free the blocks allocated for the request.
+        We free the blocks in reverse order so that the tail blocks are evicted
+        first when caching is enabled.
+
+        Args:
+            request: The request to free the blocks.
+        """
+        self.coordinator.free(request.request_id)
+
+    def remove_skipped_blocks(
+        self,
+        request_id: str,
+        processed_computed_tokens: int,
+        num_prompt_tokens: int | None = None,
+    ) -> None:
+        """Remove the blocks that are no longer needed from `blocks` and replace
+        the removed blocks with null_block.
+
+        Args:
+            request_id: The request ID.
+            processed_computed_tokens: Computed-token prefix length covering
+                fully processed and committed tokens only (safe to free).
+            num_prompt_tokens: Optional prompt length for R-SWA gap eviction.
+        """
+        self.coordinator.remove_skipped_blocks(
+            request_id, processed_computed_tokens, num_prompt_tokens
+        )
+
+    def pop_blocks_for_free(self, request: Request) -> list[KVCacheBlock]:
+        """Pop the request's bookkeeping and return its blocks without
+        returning them to the block pool. The caller must eventually free
+        them in reverse order (so that tail blocks are evicted first).
+
+        Args:
+            request: The request to pop the blocks for.
+
+        Returns:
+            The request's blocks in allocation order.
+        """
+        return self.coordinator.pop_blocks_for_free(request.request_id)
+
+    def evict_blocks(self, block_ids: set[int]) -> None:
+        """evict blocks from the prefix cache by their block IDs.
+
+        Args:
+            block_ids: Set of block IDs to evict from cache.
+        """
+        self.block_pool.evict_blocks(block_ids)
+
+    def reset_prefix_cache(self) -> bool:
+        """Reset prefix cache. This function may be used in RLHF
+        flows to invalidate prefix caching after the weights are updated,
+        or used for resetting prefix caching status for benchmarking.
+
+        Returns:
+            bool: True if the prefix cache is successfully reset,
+            False otherwise.
+        """
+        if not self.coordinator.reset_prefix_cache():
+            return False
+        if self.log_stats:
+            assert self.prefix_cache_stats is not None
+            self.prefix_cache_stats.reset = True
+        return True
+
+    def get_num_common_prefix_blocks(self, running_request_id: str) -> list[int]:
+        """Calculate the number of common prefix blocks for each kv cache group.
+
+        The function selects a running request and iterates through its blocks.
+        A block is considered a common prefix block if ALL requests with
+        allocated KV cache share it (i.e., ref_cnt equals the number of entries
+        in req_to_blocks).
+
+        NOTE(woosuk): The number of requests with allocated KV cache is **greater
+        than or equal to** the number of requests scheduled in the current step.
+        This is because having allocated KV cache only indicates that:
+        1. The request has not yet finished, and
+        2. The request holds its blocks unfreed.
+
+        While all scheduled requests must have allocated KV cache, the inverse
+        is not necessarily true. There may be requests with allocated KV cache
+        that are not scheduled in the current step.
+
+        This can result in an edge case where the number of common prefix blocks
+        is 0, even though all scheduled requests share a common prefix. This
+        occurs because there may be unscheduled requests that do not share the
+        common prefix. Currently, this case cannot be easily detected, so the
+        function returns 0 in such cases.
+
+        Args:
+            running_request_id: The request ID of any running request, used to
+                identify the common prefix blocks.
+
+        Returns:
+            list[int]: The number of common prefix blocks for each kv cache
+            group.
+        """
+        return self.coordinator.get_num_common_prefix_blocks(running_request_id)
+
+    def take_events(self) -> list[KVCacheEvent]:
+        """Take the KV cache events from the block pool.
+
+        Returns:
+            A list of KV cache events.
+        """
+        events = self.block_pool.take_events()
+        for event in events:
+            if not isinstance(event, BlockStored):
+                continue
+            if event.group_idx is None:
+                continue
+            if event.group_idx < 0 or event.group_idx >= len(
+                self.kv_cache_event_metadata
+            ):
+                logger.warning(
+                    "Group index `%s` not in KV cache metadata", event.group_idx
+                )
+                continue
+            # Annotate here so BlockPool can keep emitting structural cache
+            # events without owning semantic KV cache spec metadata.
+            kind, sliding_window = self.kv_cache_event_metadata[event.group_idx]
+            event.kv_cache_spec_kind = kind
+            event.kv_cache_spec_sliding_window = sliding_window
+        return events
+
+    def get_blocks(self, request_id: str) -> KVCacheBlocks:
+        """Get the blocks of a request."""
+        return self.create_kv_cache_blocks(self.coordinator.get_blocks(request_id))
+
+    def get_block_ids(self, request_id: str) -> tuple[list[int], ...]:
+        """Get the block ids of a request."""
+        return self.get_blocks(request_id).get_block_ids()
+
+    def get_block_ids_for_computed_tokens(
+        self,
+        request_id: str,
+        num_computed_tokens: int,
+    ) -> tuple[list[int], ...]:
+        """Get block ids covering the request's computed tokens."""
+        block_ids = self.get_block_ids(request_id)
+        clipped_block_ids: list[list[int]] = []
+        for group, ids in zip(self.kv_cache_config.kv_cache_groups, block_ids):
+            spec = group.kv_cache_spec
+            if not isinstance(spec, AttentionSpec) or isinstance(
+                spec, (CrossAttentionSpec, EncoderOnlyAttentionSpec)
+            ):
+                clipped_block_ids.append(ids)
+                continue
+
+            num_valid_blocks = cdiv(num_computed_tokens, spec.block_size)
+            clipped_block_ids.append(ids[:num_valid_blocks])
+        return tuple(clipped_block_ids)
+
+    def estimate_cached_tokens(self, request: Request) -> int:
+        """Estimate the number of tokens cached by the request."""
+        cached_tokens: int | None = None
+        for group, blocks in zip(
+            self.kv_cache_config.kv_cache_groups,
+            self.get_blocks(request.request_id).blocks,
+        ):
+            if not group.kv_cache_spec.prefix_cacheable:
+                continue
+            if isinstance(
+                group.kv_cache_spec,
+                (CrossAttentionSpec, EncoderOnlyAttentionSpec),
+            ):
+                # Cross-attention and encoder-only groups are not prefix cached.
+                continue
+
+            group_cached_tokens = 0
+            for block in blocks:
+                group_cached_tokens = max(
+                    group_cached_tokens,
+                    block.block_hash_num_tokens or 0,
+                )
+
+            cached_tokens = (
+                group_cached_tokens
+                if cached_tokens is None
+                else min(cached_tokens, group_cached_tokens)
+            )
+
+        return cached_tokens or 0
+
+    def cache_blocks(self, request: Request, num_computed_tokens: int) -> None:
+        """Cache the blocks for the request, if enabled.
+
+        Args:
+            request: The request to cache the blocks.
+            num_computed_tokens: The number of computed tokens, including tokens
+                that are already cached and tokens to be cached.
+        """
+        if self.enable_caching:
+            self.coordinator.cache_blocks(request, num_computed_tokens)
+
+    def create_kv_cache_blocks(
+        self, blocks: tuple[list[KVCacheBlock], ...]
+    ) -> KVCacheBlocks:
+        # Only create new KVCacheBlocks for non-empty blocks
+        return KVCacheBlocks(blocks) if any(blocks) else self.empty_kv_cache_blocks
+
+    def truncate_computed_blocks(
+        self, blocks: KVCacheBlocks, num_computed_tokens: int
+    ) -> KVCacheBlocks:
+        """Return a lookup-result view truncated at an aligned token endpoint.
+
+        An external hit can supply the final Mamba state even when the local
+        Mamba group ends before this endpoint. Other groups must cover it.
+        Pure slicing: refcounts are untouched and ``blocks`` is not mutated.
+        """
+        truncated: list[list[KVCacheBlock]] = []
+        for group_blocks, manager, group in zip(
+            blocks.blocks,
+            self.coordinator.single_type_managers,
+            self.kv_cache_config.kv_cache_groups,
+            strict=True,
+        ):
+            if not group.kv_cache_spec.prefix_cacheable:
+                # Scratch groups (e.g. the CSA compressor ring) hold a fixed
+                # block covering no token range, so a lookup result never
+                # carries computed blocks for them and their block size does
+                # not divide the endpoint.
+                assert not group_blocks
+                truncated.append([])
+                continue
+            assert num_computed_tokens % manager.block_size == 0
+            num_blocks = num_computed_tokens // manager.block_size
+            if isinstance(group.kv_cache_spec, MambaSpec):
+                num_blocks = min(num_blocks, len(group_blocks))
+            else:
+                assert num_blocks <= len(group_blocks)
+            truncated.append(list(group_blocks[:num_blocks]))
+        return self.create_kv_cache_blocks(tuple(truncated))
+
+    def take_new_block_ids(self) -> list[int]:
+        """Drain and return new attention block IDs for zeroing."""
+        ids: list[int] = []
+        for mgr in self.coordinator.single_type_managers:
+            ids.extend(mgr.take_new_block_ids())
+        return ids
+
+    def get_zeroing_block_ids_in_range(
+        self, request_id: str, start_token: int, end_token: int
+    ) -> list[int]:
+        """The request's block ids covering [start_token, end_token), from
+        the groups whose new blocks are zeroed by the worker."""
+        ids: list[int] = []
+        for mgr in self.coordinator.single_type_managers:
+            if mgr.records_new_block_ids:
+                start_idx = start_token // mgr.block_size
+                end_idx = cdiv(end_token, mgr.block_size)
+                blocks = mgr.req_to_blocks[request_id]
+                ids.extend(blk.block_id for blk in blocks[start_idx:end_idx])
+        return ids
+
+    def record_blocks_for_zeroing(self, request_id: str, start_token: int) -> None:
+        """Re-record the request's blocks from start_token onwards for
+        zeroing, e.g. blocks a failed async KV load left unwritten.
+
+        start_token must be block-aligned: zeroing a partially-valid block
+        would wipe its valid prefix.
+        """
+        for mgr in self.coordinator.single_type_managers:
+            if mgr.records_new_block_ids:
+                assert start_token % mgr.block_size == 0
+                start_idx = start_token // mgr.block_size
+                blocks = mgr.req_to_blocks[request_id]
+                mgr.new_block_ids.extend(blk.block_id for blk in blocks[start_idx:])
+
+    def take_kv_cache_block_copies(
+        self,
+    ) -> tuple[list[KVCacheBlockCopy], list[KVCacheBlock]]:
+        """Drain pending copies and return their retained endpoints."""
+        pending_copies: list[tuple[KVCacheBlock, KVCacheBlock]] = []
+        for mgr in self.coordinator.single_type_managers:
+            pending_copies.extend(mgr.take_pending_cow_copies())
+        copies = [
+            KVCacheBlockCopy(
+                src_block_id=source_block.block_id,
+                dst_block_id=cow_block.block_id,
+            )
+            for source_block, cow_block in pending_copies
+        ]
+        retained_blocks = [block for pair in pending_copies for block in pair]
+        return copies, retained_blocks
+
+    def take_boundary_state_offloads(
+        self,
+    ) -> dict[str, list[tuple[int, int, int]]]:
+        """Drain this step's boundary-state hand-offs for a KV connector.
+
+        Returns ``{request_id: [(group_id, block_id, boundary_tokens), ...]}``
+        for mamba "align" boundary states: the
+        request's committed boundary-state snapshots and, on a sub-block partial
+        hit, the CoW copy of its last-prompt-boundary state. Only mamba "align"
+        groups contribute; empty otherwise. A connector reads the referenced
+        blocks — never resolving them positionally — and offloads them so a
+        later request can hit that prefix.
+        """
+        offloads: dict[str, list[tuple[int, int, int]]] = {}
+        for mgr in self.coordinator.single_type_managers:
+            for (
+                req_id,
+                group_id,
+                block,
+                boundary_tokens,
+            ) in mgr.take_pending_boundary_state_offloads():
+                offloads.setdefault(req_id, []).append(
+                    (group_id, block.block_id, boundary_tokens)
+                )
+        return offloads
+
+    def finalize_partial_tail_offloads(
+        self, request: Request
+    ) -> list[tuple[int, int, int]]:
+        """Consume safe producer partial tails when a request finishes.
+
+        A mamba align table block is a valid boundary source only if no later
+        token was forwarded. The connector pins and queues the exact table
+        block before request cleanup, then releases the pin when every worker
+        reports the store job complete.
+        """
+        offloads: list[tuple[int, int, int]] = []
+        for mgr in self.coordinator.single_type_managers:
+            finalized = mgr.finalize_partial_tail_offload(
+                request.request_id,
+                request.num_computed_tokens,
+                request.num_in_flight_tokens,
+            )
+            if finalized is None:
+                continue
+            group_id, block, boundary_tokens = finalized
+            offloads.append((group_id, block.block_id, boundary_tokens))
+        return offloads
+
+    def new_step_starts(self) -> None:
+        """Notify the coordinator that a new step is starting."""
+        self.coordinator.new_step_starts()

serving/fork/v1/core/sched/output.py

@@ -0,0 +1,336 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from collections.abc import Callable
+from dataclasses import dataclass, field
+from functools import cached_property
+from typing import TYPE_CHECKING
+
+from vllm.config.ec_manager_config import EncoderCacheManagerMetadata
+from vllm.multimodal.utils import strip_covered_mm_data
+
+if TYPE_CHECKING:
+    import numpy as np
+    import numpy.typing as npt
+    import torch
+
+    from vllm.distributed.ec_transfer.ec_connector.base import ECConnectorMetadata
+    from vllm.distributed.kv_transfer.kv_connector.v1.base import KVConnectorMetadata
+    from vllm.lora.request import LoRARequest
+    from vllm.multimodal.inputs import MultiModalFeatureSpec
+    from vllm.pooling_params import PoolingParams
+    from vllm.sampling_params import SamplingParams
+    from vllm.v1.core.kv_cache_utils import KVCacheBlockCopy
+    from vllm.v1.request import Request
+else:
+    ECConnectorMetadata = object
+    KVConnectorMetadata = object
+    KVCacheBlockCopy = object
+    LoRARequest = object
+    MultiModalFeatureSpec = object
+    PoolingParams = object
+    SamplingParams = object
+    Request = object
+
+
+@dataclass
+class NewRequestData:
+    req_id: str
+    prompt_token_ids: list[int] | None
+    mm_features: list[MultiModalFeatureSpec]
+    sampling_params: SamplingParams | None
+    pooling_params: PoolingParams | None
+    block_ids: tuple[list[int], ...]
+    num_computed_tokens: int
+    lora_request: LoRARequest | None
+    prompt_embeds: "torch.Tensor | None" = None
+    prompt_is_token_ids: list[bool] | None = None
+
+    # Only used for v2 model runner.
+    prefill_token_ids: list[int] | None = None
+
+    @classmethod
+    def from_request(
+        cls,
+        request: Request,
+        block_ids: tuple[list[int], ...],
+        prefill_token_ids: list[int] | None = None,
+        uses_mrope: bool = False,
+    ) -> "NewRequestData":
+        return cls(
+            req_id=request.request_id,
+            prompt_token_ids=request.prompt_token_ids,
+            mm_features=strip_covered_mm_data(
+                request.mm_features,
+                request.num_computed_tokens,
+                uses_mrope=uses_mrope,
+            ),
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            block_ids=block_ids,
+            num_computed_tokens=request.num_computed_tokens,
+            lora_request=request.lora_request,
+            prompt_embeds=request.prompt_embeds,
+            prompt_is_token_ids=request.prompt_is_token_ids,
+            prefill_token_ids=prefill_token_ids,
+        )
+
+    @property
+    def prompt_len(self) -> int:
+        if self.prompt_token_ids is not None:
+            return len(self.prompt_token_ids)
+        if self.prompt_embeds is not None:
+            return self.prompt_embeds.shape[0]
+        return 0
+
+    def __repr__(self) -> str:
+        prompt_embeds_shape = (
+            self.prompt_embeds.shape if self.prompt_embeds is not None else None
+        )
+        return (
+            f"NewRequestData("
+            f"req_id={self.req_id},"
+            f"prompt_token_ids={self.prompt_token_ids},"
+            f"prefill_token_ids={self.prefill_token_ids},"
+            f"mm_features={self.mm_features},"
+            f"sampling_params={self.sampling_params},"
+            f"block_ids={self.block_ids},"
+            f"num_computed_tokens={self.num_computed_tokens},"
+            f"lora_request={self.lora_request},"
+            f"prompt_embeds_shape={prompt_embeds_shape}"
+            ")"
+        )
+
+    # Version of __repr__ with the prompt data obfuscated
+    def anon_repr(self) -> str:
+        prompt_token_ids_len = (
+            len(self.prompt_token_ids) if self.prompt_token_ids is not None else None
+        )
+        prompt_embeds_shape = (
+            self.prompt_embeds.shape if self.prompt_embeds is not None else None
+        )
+        prefill_token_ids_len = (
+            len(self.prefill_token_ids) if self.prefill_token_ids is not None else None
+        )
+        return (
+            f"NewRequestData("
+            f"req_id={self.req_id},"
+            f"prompt_token_ids_len={prompt_token_ids_len},"
+            f"prefill_token_ids_len={prefill_token_ids_len},"
+            f"mm_features={self.mm_features},"
+            f"sampling_params={self.sampling_params},"
+            f"block_ids={self.block_ids},"
+            f"num_computed_tokens={self.num_computed_tokens},"
+            f"lora_request={self.lora_request},"
+            f"prompt_embeds_shape={prompt_embeds_shape}"
+            ")"
+        )
+
+
+@dataclass
+class CachedRequestData:
+    req_ids: list[str]
+    # For request ids not in resumed_req_ids, new_block_ids will be appended to
+    # the request's block IDs. For those in the set, new_block_ids will be used as the
+    # request's block IDs instead of appending to the existing block IDs.
+    resumed_req_ids: set[str]
+    # NOTE(woosuk): new_token_ids is only used for pipeline parallelism.
+    # When PP is not used, new_token_ids will be empty.
+    new_token_ids: list[list[int]]
+    # MRV1-only: For requests not scheduled in the last step, propagate the token ids
+    # to the connector. Won't contain requests scheduled in the prior step.
+    all_token_ids: dict[str, list[int]]
+    new_block_ids: list[tuple[list[int], ...] | None]
+    num_computed_tokens: list[int]
+    num_output_tokens: list[int]
+
+    # Version of dataclass repr with token IDs obfuscated.
+    def anon_repr(self) -> str:
+        new_token_ids_lens = [len(toks) for toks in self.new_token_ids]
+        all_token_ids_lens = {
+            req_id: len(toks) for req_id, toks in self.all_token_ids.items()
+        }
+        return (
+            f"CachedRequestData("
+            f"req_ids={self.req_ids},"
+            f"resumed_req_ids={self.resumed_req_ids},"
+            f"new_token_ids_lens={new_token_ids_lens},"
+            f"all_token_ids_lens={all_token_ids_lens},"
+            f"new_block_ids={self.new_block_ids},"
+            f"num_computed_tokens={self.num_computed_tokens},"
+            f"num_output_tokens={self.num_output_tokens}"
+            f")"
+        )
+
+    def __repr__(self) -> str:
+        return self.anon_repr()
+
+    @property
+    def num_reqs(self) -> int:
+        return len(self.req_ids)
+
+    @cached_property
+    def _req_id_to_num_output_tokens(self) -> dict[str, int]:
+        """Cache mapping of req_id to num_output_tokens for O(1) lookup.
+
+        This cached property is safe because CachedRequestData instances
+        are created fresh each scheduling iteration and not mutated during
+        computation of iteration details.
+        """
+        return dict(zip(self.req_ids, self.num_output_tokens))
+
+    def is_context_phase(self, req_id: str) -> bool:
+        num_output_tokens = self._req_id_to_num_output_tokens.get(req_id)
+        return num_output_tokens is not None and num_output_tokens == 0
+
+    @classmethod
+    def make_empty(cls) -> "CachedRequestData":
+        return cls(
+            req_ids=[],
+            resumed_req_ids=set(),
+            new_token_ids=[],
+            all_token_ids={},
+            new_block_ids=[],
+            num_computed_tokens=[],
+            num_output_tokens=[],
+        )
+
+
+@dataclass
+class ScheduledEncoderInputStats:
+    """Stats for encoder inputs scheduled in one iteration."""
+
+    num_inputs: int = 0
+    output_tokens: int = 0
+
+
+@dataclass
+class KVConnectorBlockState:
+    """Scheduler-local block state offered to a producer-side KV connector."""
+
+    # Requests scheduled this step and requests with a boundary-state hand-off.
+    req_ids: set[str]
+    # Resolve on access to avoid copying tables the connector never reads.
+    resolve_block_ids: Callable[[str], tuple[list[int], ...]]
+    # Exact Mamba "align" boundary-state hand-offs.
+    boundary_state_offloads: dict[str, list[tuple[int, int, int]]]
+
+    def get_block_ids(self, req_id: str) -> tuple[list[int], ...] | None:
+        if req_id not in self.req_ids:
+            return None
+        return self.resolve_block_ids(req_id)
+
+
+@dataclass
+class SchedulerOutput:
+    # list of the requests that are scheduled for the first time.
+    # We cache the request's data in each worker process, so that we don't
+    # need to re-send it every scheduling step.
+    scheduled_new_reqs: list[NewRequestData]
+    # list of the requests that have been scheduled before.
+    # Since the request's data is already cached in the worker processes,
+    # we only send the diff to minimize the communication cost.
+    scheduled_cached_reqs: CachedRequestData
+
+    # req_id -> num_scheduled_tokens
+    # Number of tokens scheduled for each request.
+    num_scheduled_tokens: dict[str, int]
+    # Total number of tokens scheduled for all requests.
+    # Equal to sum(num_scheduled_tokens.values())
+    total_num_scheduled_tokens: int
+    # req_id -> spec_token_ids
+    # If a request does not have any spec decode tokens, it will not be
+    # included in the dictionary.
+    scheduled_spec_decode_tokens: dict[str, list[int]]
+    # req_id -> encoder input indices that need processing.
+    # E.g., if a request has [0, 1], it could mean the vision encoder needs
+    # to process that the request's 0-th and 1-th images in the current step.
+    scheduled_encoder_inputs: dict[str, list[int]]
+    # Number of common prefix blocks for all requests in each KV cache group.
+    # This can be used for cascade attention.
+    num_common_prefix_blocks: list[int]
+
+    # Request IDs that are finished in between the previous and the current
+    # steps. This is used to notify the workers about the finished requests
+    # so that they can free the cached states for those requests.
+    finished_req_ids: set[str]
+    # list of mm_hash strings associated with the encoder outputs to be
+    # freed from the encoder cache.
+    free_encoder_mm_hashes: list[str]
+    # FORK: req_id -> number of zero-token latent KV positions scheduled this
+    # step (HybridAgent latent steps consume a KV position without a token id).
+    # Empty when no latent step is scheduled.
+    latent_steps: dict[str, int] = field(default_factory=dict)
+    # FORK: req_id -> 1 for the latent ENTRY step: a pending-latent request
+    # that still has a pending token id (the `<latent>` token position). The
+    # model runner drops that step's sample so the position becomes a latent
+    # KV position fed by project_latent(prev_hidden).
+    latent_entry_steps: dict[str, int] = field(default_factory=dict)
+
+    scheduled_encoder_input_stats: ScheduledEncoderInputStats | None = None
+
+    # Request IDs that are preempted in this step.
+    # Only used for v2 model runner.
+    preempted_req_ids: set[str] | None = None
+
+    # Whether any of the scheduled requests use structured output.
+    # Set only in async scheduling case.
+    has_structured_output_requests: bool = False
+
+    # Whether the scheduled requests have all the output tokens they
+    # need to perform grammar bitmask computation.
+    pending_structured_output_tokens: bool = False
+
+    # Used for adjusting acceptance rate calculation.
+    num_invalid_spec_tokens: dict[str, int] | None = None
+
+    # KV Cache Connector metadata.
+    kv_connector_metadata: KVConnectorMetadata | None = None
+
+    # Whether any scheduled request consumes KV that the connector loads
+    # synchronously during this step (load_async=False).
+    has_sync_kv_loads: bool = False
+
+    # EC Cache Connector metadata
+    ec_connector_metadata: ECConnectorMetadata | None = None
+    # EC Cache Manager metadata
+    ec_manager_metadata: EncoderCacheManagerMetadata | None = None
+    # Block IDs freshly allocated from the pool during this scheduling step.
+    # The worker zeros the corresponding GPU memory before the blocks are used,
+    # preventing stale NaN/data from corrupting attention or SSM computation.
+    new_block_ids_to_zero: list[int] | None = None
+
+    # CoW copies to apply after zeroing new blocks and before forward.
+    kv_cache_block_copies: list[KVCacheBlockCopy] | None = None
+
+    # Complete block-table rows that replace incrementally appended block IDs.
+    block_table_updates: dict[str, tuple[list[int], ...]] | None = None
+
+    # Scheduler-local; always None by the time this reaches a worker.
+    kv_connector_block_state: KVConnectorBlockState | None = None
+
+    # Dynamic speculative decoding: optimal K chosen by scheduler.
+    # Number of spec tokens to schedule for the next step.
+    num_spec_tokens_to_schedule: int = 0
+
+    @classmethod
+    def make_empty(cls) -> "SchedulerOutput":
+        return cls(
+            scheduled_new_reqs=[],
+            scheduled_cached_reqs=CachedRequestData.make_empty(),
+            num_scheduled_tokens={},
+            total_num_scheduled_tokens=0,
+            scheduled_spec_decode_tokens={},
+            scheduled_encoder_inputs={},
+            num_common_prefix_blocks=[],
+            finished_req_ids=set(),
+            free_encoder_mm_hashes=[],
+        )
+
+
+@dataclass
+class GrammarOutput:
+    # ids of structured output requests.
+    structured_output_request_ids: list[str]
+    # Bitmask ordered as structured_output_request_ids.
+    grammar_bitmask: "npt.NDArray[np.int32]"

serving/fork/v1/request.py

@@ -0,0 +1,417 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import enum
+import time
+from collections import deque
+from collections.abc import Callable, Mapping
+from dataclasses import dataclass
+from typing import TYPE_CHECKING, Any
+
+import torch
+
+from vllm.multimodal.inputs import MultiModalFeatureSpec
+from vllm.pooling_params import PoolingParams
+from vllm.sampling_params import SamplingParams
+from vllm.utils import length_from_prompt_token_ids_or_embeds
+from vllm.v1.engine import (
+    EngineCoreEvent,
+    EngineCoreEventType,
+    EngineCoreRequest,
+    FinishReason,
+)
+from vllm.v1.metrics.stats import PrefillStats, RequestSpecDecodeMetrics
+from vllm.v1.structured_output.request import StructuredOutputRequest
+from vllm.v1.utils import ConstantList
+
+if TYPE_CHECKING:
+    from vllm.lora.request import LoRARequest
+    from vllm.v1.core.kv_cache_utils import BlockHash
+
+
+@dataclass
+class StreamingUpdate:
+    """Lightweight data for streaming session continuation.
+
+    Contains only the fields needed to update an existing streaming session
+    with new input data.
+    """
+
+    mm_features: list[MultiModalFeatureSpec] | None
+    prompt_token_ids: list[int] | None
+    max_tokens: int
+    arrival_time: float
+    sampling_params: SamplingParams | None
+
+    @classmethod
+    def from_request(cls, request: "Request") -> "StreamingUpdate | None":
+        if not request.resumable:
+            return None
+        return cls(
+            mm_features=request.mm_features,
+            prompt_token_ids=request.prompt_token_ids,
+            max_tokens=request.max_tokens,
+            arrival_time=request.arrival_time,
+            sampling_params=request.sampling_params,
+        )
+
+
+class Request:
+    def __init__(
+        self,
+        request_id: str,
+        prompt_token_ids: list[int] | None,
+        sampling_params: SamplingParams | None,
+        pooling_params: PoolingParams | None,
+        client_index: int = 0,
+        arrival_time: float | None = None,
+        prompt_embeds: torch.Tensor | None = None,
+        prompt_is_token_ids: list[bool] | None = None,
+        mm_features: list[MultiModalFeatureSpec] | None = None,
+        lora_request: "LoRARequest | None" = None,
+        cache_salt: str | None = None,
+        priority: int = 0,
+        trace_headers: Mapping[str, str] | None = None,
+        block_hasher: Callable[["Request"], list["BlockHash"]] | None = None,
+        resumable: bool = False,
+        session_id: str | None = None,
+        reasoning_ended: bool | None = None,
+        reasoning_parser_kwargs: dict[str, Any] | None = None,
+        abort_immediately: bool = False,
+    ) -> None:
+        self.request_id = request_id
+        self.client_index = client_index
+        self.priority = priority
+        self.sampling_params = sampling_params
+        self.pooling_params = pooling_params
+        self.lora_request = lora_request
+        self.structured_output_request = StructuredOutputRequest.from_sampling_params(
+            sampling_params
+        )
+        if self.structured_output_request is not None:
+            self.structured_output_request.reasoning_ended = reasoning_ended
+            self.structured_output_request.reasoning_parser_kwargs = (
+                reasoning_parser_kwargs
+            )
+        self.arrival_time = arrival_time if arrival_time is not None else time.time()
+
+        self.status = RequestStatus.WAITING
+        self.events: list[EngineCoreEvent] = []
+        self.stop_reason: int | str | None = None
+
+        # P/D: Connector-specific KV transfer parameters.
+        self.kv_transfer_params: dict[str, Any] | None = None
+        # E/P/D: Connector-specific encoder-cache transfer parameters.
+        self.ec_transfer_params: dict[str, Any] | None = None
+
+        if pooling_params is not None:
+            # Pooling models.
+            self.max_tokens = 1
+        elif sampling_params is not None:
+            # Generative models.
+            assert sampling_params.max_tokens is not None
+            self.max_tokens = sampling_params.max_tokens
+            if self.structured_output_request is not None:
+                self.status = RequestStatus.WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR
+
+            if sampling_params.extra_args is not None:
+                self.kv_transfer_params = sampling_params.extra_args.get(
+                    "kv_transfer_params"
+                )
+                self.ec_transfer_params = sampling_params.extra_args.get(
+                    "ec_transfer_params"
+                )
+                self.kv_cache_report_mode = sampling_params.extra_args.get(
+                    "kv_cache_report_mode", "incremental"
+                )
+            else:
+                self.kv_cache_report_mode = "incremental"
+        else:
+            raise ValueError("sampling_params and pooling_params can't both be unset")
+
+        self.prompt_token_ids = prompt_token_ids
+        self.prompt_embeds = prompt_embeds
+        # Per-position mask used in mixed-mode (chat completion with
+        # prompt_embeds). `None` except when both `prompt_token_ids` and
+        # `prompt_embeds` are set and their positions are interleaved.
+        self.prompt_is_token_ids = prompt_is_token_ids
+        # Cache per-block prompt-embed hashes to avoid rehashing the same
+        # tensor slices when generating extra keys.
+        self._prompt_embeds_per_block_hashes: dict[tuple[int, int], bytes] = {}
+        self.num_prompt_tokens = length_from_prompt_token_ids_or_embeds(
+            prompt_token_ids, prompt_embeds
+        )
+        self._output_token_ids: list[int] = []
+        if self.prompt_token_ids is None:
+            self._all_token_ids: list[int] = [0] * self.num_prompt_tokens
+        elif self.prompt_is_token_ids is None:
+            self._all_token_ids = self.prompt_token_ids.copy()
+        else:
+            # Mixed-mode prompt: positions covered by prompt_embeds hold a sentinel
+            # special token id that may lie outside the embedding. Zero them, matching
+            # the no-token-ids case above, so embedding gathers over these placeholder
+            # ids stay in bounds; the actual inputs come from prompt_embeds.
+            self._all_token_ids = [
+                t if is_tok else 0
+                for t, is_tok in zip(self.prompt_token_ids, self.prompt_is_token_ids)
+            ]
+
+        # Used in async scheduling.
+        self.num_output_placeholders = 0
+        # Tokens of output in flight when the request was preempted: delivered
+        # on return, but must not mutate the reset counters.
+        self.num_stale_output_tokens = 0
+        # Drop the stale output instead, for same-step preempt + resume
+        # (reset_prefix_cache).
+        self.drop_stale_output = False
+
+        # Tokens of steps whose output is not yet processed (async scheduling
+        # and PP run ahead of the GPU); `num_computed_tokens` counts them
+        # optimistically.
+        self.num_in_flight_tokens = 0
+
+        # V2+PP+async: Enforces `pp_size` cadence between same-request decode steps
+        # so the worker's broadcast slot ring stays consistent.
+        self.next_decode_eligible_step = 0
+
+        # Seq of the most recent step this request was scheduled in; fences
+        # deferred block freeing (see Scheduler._free_request_blocks).
+        self.last_sched_seq = 0
+
+        self.spec_token_ids: list[int] = []
+        self.num_computed_tokens = 0
+        # FORK: zero-token latent steps (HybridAgent). Latent KV positions
+        # filled so far; invariant:
+        # num_computed_tokens == real token KV positions + num_latent_computed.
+        self.num_latent_computed = 0
+        # FORK: frozen to num_computed_tokens at the first latent step; latent
+        # blocks are allocated but never prefix-cached (their KV derives from
+        # an injected embedding, not a token id).
+        self.kv_cacheable_prefix: int | None = None
+        # FORK: set by the state machine (T4/T5) to mark the next step as a
+        # zero-token latent step.
+        self.pending_latent = False
+        self.cache_salt: str | None = cache_salt
+
+        # Multi-modal related
+        self.mm_features = mm_features or []
+
+        # Read-only views
+        # Prevent directly appending to these lists since
+        # they should also be updated simultaneously.
+        self.output_token_ids = ConstantList(self._output_token_ids)
+        self.all_token_ids = ConstantList(self._all_token_ids)
+        # trace_headers
+        self.trace_headers = trace_headers
+        self.session_id = session_id
+
+        # True if this request is scheduled as a non-final prefill chunk.
+        self.is_prefill_chunk = False
+
+        # Block-aligned token position of a proven shared prefix worth pinning
+        # in the (sparse) prefix cache; 0 means none. Set at admission for
+        # hybrid/Mamba models when a shared prefix is detected (Marconi-style).
+        self.shared_prefix_boundary = 0
+
+        # The number of NaNs in logits. A value greater than 0
+        # indicates that the output is corrupted
+        self.num_nans_in_logits = 0
+
+        # The number of times this request has been preempted by the scheduler.
+        self.num_preemptions = 0
+
+        self.prefill_stats: PrefillStats | None = PrefillStats()
+
+        # Per-request speculative-decoding acceptance accumulator. Populated by
+        # the scheduler when --per-request-spec-decode-metrics is set (eagerly on
+        # add_request, then observed each verify step); stays None otherwise.
+        self.spec_decode_metrics: RequestSpecDecodeMetrics | None = None
+
+        self.block_hashes: list[BlockHash] = []
+        # Store the block hasher without binding self to avoid creating a
+        # reference cycle (Request -> partial -> Request) that prevents
+        # immediate garbage collection via reference counting.
+        self._block_hasher: Callable[[Request], list[BlockHash]] | None = block_hasher
+        self.update_block_hashes()
+
+        self.skip_reading_prefix_cache = self.get_skip_reading_prefix_cache()
+
+        # Used for streaming
+        self.resumable = resumable
+        # None entry in the queue means finished.
+        self.streaming_queue: deque[StreamingUpdate | None] | None = None
+
+        # If True, request should be aborted immediately after being added to
+        # the scheduler so the connector's request_finished hook runs.
+        self.abort_immediately = abort_immediately
+
+    @classmethod
+    def from_engine_core_request(
+        cls,
+        request: EngineCoreRequest,
+        block_hasher: Callable[["Request"], list["BlockHash"]] | None,
+    ) -> "Request":
+        return cls(
+            request_id=request.request_id,
+            client_index=request.client_index,
+            prompt_token_ids=request.prompt_token_ids,
+            prompt_embeds=request.prompt_embeds,
+            prompt_is_token_ids=request.prompt_is_token_ids,
+            mm_features=request.mm_features,
+            sampling_params=request.sampling_params,
+            pooling_params=request.pooling_params,
+            arrival_time=request.arrival_time,
+            lora_request=request.lora_request,
+            cache_salt=request.cache_salt,
+            priority=request.priority,
+            trace_headers=request.trace_headers,
+            block_hasher=block_hasher,
+            resumable=request.resumable,
+            session_id=request.session_id,
+            reasoning_ended=request.reasoning_ended,
+            reasoning_parser_kwargs=request.reasoning_parser_kwargs,
+            abort_immediately=request.abort_immediately,
+        )
+
+    def append_output_token_ids(
+        self,
+        token_ids: int | list[int],
+    ) -> None:
+        if isinstance(token_ids, int):
+            self._output_token_ids.append(token_ids)
+            self._all_token_ids.append(token_ids)
+        else:
+            self._output_token_ids.extend(token_ids)
+            self._all_token_ids.extend(token_ids)
+
+        self.update_block_hashes()
+
+    def update_block_hashes(self) -> None:
+        """Compute block hashes for any new full blocks and append them."""
+        if self._block_hasher is not None:
+            self.block_hashes.extend(self._block_hasher(self))
+
+    @property
+    def use_structured_output(self) -> bool:
+        return self.structured_output_request is not None
+
+    @property
+    def num_tokens(self) -> int:
+        return len(self._all_token_ids)
+
+    @property
+    def num_tokens_with_spec(self) -> int:
+        return len(self._all_token_ids) + len(self.spec_token_ids)
+
+    @property
+    def num_output_tokens(self) -> int:
+        return len(self._output_token_ids)
+
+    @property
+    def num_encoder_inputs(self) -> int:
+        return len(self.mm_features)
+
+    @property
+    def has_encoder_inputs(self) -> bool:
+        return self.num_encoder_inputs > 0
+
+    def get_skip_reading_prefix_cache(self) -> bool:
+        if (
+            self.sampling_params is not None
+            and self.sampling_params.skip_reading_prefix_cache is not None
+        ):
+            return self.sampling_params.skip_reading_prefix_cache
+        elif (
+            self.pooling_params is not None
+            and self.pooling_params.skip_reading_prefix_cache is not None
+        ):
+            return self.pooling_params.skip_reading_prefix_cache
+        return False
+
+    def is_finished(self) -> bool:
+        return RequestStatus.is_finished(self.status)
+
+    def get_finished_reason(self) -> FinishReason | None:
+        return RequestStatus.get_finished_reason(self.status)
+
+    def get_num_encoder_embeds(self, input_id: int) -> int:
+        assert input_id < len(self.mm_features)
+        return self.mm_features[input_id].mm_position.get_num_embeds()
+
+    def record_event(
+        self,
+        event_type: EngineCoreEventType,
+        timestamp: float | None = None,
+    ) -> None:
+        self.events.append(EngineCoreEvent.new_event(event_type, timestamp))
+
+    def take_events(self) -> list[EngineCoreEvent] | None:
+        if not self.events:
+            return None
+        events, self.events = self.events, []
+        return events
+
+    def take_prefill_stats(self) -> PrefillStats | None:
+        if self.prefill_stats is None:
+            return None
+        prefill_stats = self.prefill_stats
+        self.prefill_stats = None
+        return prefill_stats
+
+    def __lt__(self, other: "Request") -> bool:
+        """
+        Compare two requests based on priority, arrival time, and request ID.
+        Used in priority scheduling.
+        """
+        if self.priority != other.priority:
+            return self.priority < other.priority
+        if self.arrival_time != other.arrival_time:
+            return self.arrival_time < other.arrival_time
+        if self.request_id != other.request_id:
+            return self.request_id < other.request_id
+        return id(self) < id(other)
+
+
+class RequestStatus(enum.IntEnum):
+    """Status of a request."""
+
+    WAITING = enum.auto()
+    WAITING_FOR_STRUCTURED_OUTPUT_GRAMMAR = enum.auto()
+    WAITING_FOR_REMOTE_KVS = enum.auto()
+    WAITING_FOR_STREAMING_REQ = enum.auto()
+    RUNNING = enum.auto()
+    PREEMPTED = enum.auto()
+    # Note: anything after PREEMPTED will be considered
+    # as a finished status.
+    FINISHED_STOPPED = enum.auto()
+    FINISHED_LENGTH_CAPPED = enum.auto()
+    FINISHED_ABORTED = enum.auto()
+    FINISHED_IGNORED = enum.auto()
+    FINISHED_ERROR = enum.auto()
+    FINISHED_REPETITION = enum.auto()
+
+    def __str__(self) -> str:
+        return self.name
+
+    @staticmethod
+    def is_finished(status: "RequestStatus") -> bool:
+        return status > RequestStatus.PREEMPTED
+
+    @staticmethod
+    def get_finished_reason(status: "RequestStatus") -> FinishReason | None:
+        return _FINISHED_REASON_MAP.get(status)
+
+
+# Mapping of finished statuses to their finish reasons.
+# NOTE: The ignored requests are the requests whose prompt lengths
+# are longer than the model's length cap. Therefore, the stop
+# reason should also be "length" as in OpenAI API.
+_FINISHED_REASON_MAP = {
+    RequestStatus.FINISHED_STOPPED: FinishReason.STOP,
+    RequestStatus.FINISHED_LENGTH_CAPPED: FinishReason.LENGTH,
+    RequestStatus.FINISHED_ABORTED: FinishReason.ABORT,
+    RequestStatus.FINISHED_IGNORED: FinishReason.LENGTH,
+    RequestStatus.FINISHED_ERROR: FinishReason.ERROR,
+    RequestStatus.WAITING_FOR_STREAMING_REQ: FinishReason.STOP,
+    RequestStatus.FINISHED_REPETITION: FinishReason.REPETITION,
+}

serving/orchestrator/README.md

@@ -0,0 +1,58 @@
+# hybridagent-orchestrator (T5)
+
+The EXTERNAL ORCHESTRATOR for the forked vLLM deployment: it wraps an
+IN-PROCESS `vllm.LLMEngine` and runs the HybridAgent agentic loop that
+the T4 state machine (`processor/`) cannot finish on its own.
+
+## What it does
+
+- **Engine construction** with the fork's required semantics:
+  `VLLM_USE_V2_MODEL_RUNNER=0` (the T2 fork modifies the V1 runner; the
+  V2 default runner silently disables the latent path) and
+  `async_scheduling=False` (the async scheduler keeps `num_new_tokens > 0`
+  every step, which blocks the fork's zero-token latent gate).
+- **Arms the T4 state machine** via `processor.register_engine(...)` with
+  `get_request` / `get_req_ids` closures bound to the in-process
+  `scheduler.requests` and `model_runner.input_batch.req_ids`, the
+  tokenizer, a `tool_executor` and the in-process `inject_token`.
+- **`add_request(...)`** with the `<latent>`-capable sampling params
+  (greedy, no penalties/structured output, a `max_tokens` budget).
+- **`run(...)`** step loop: `engine.step()`; after each step, for every
+  request paused in `IDLE`, read `pending_tool` (tag, content), call
+  `tool_executor(tag, content)`, tokenize `<result>{...}</result>`
+  (`add_special_tokens=False`) and `queue_observation(rid, tokens)`;
+  collect the final text on DONE.
+- **Entry-step alignment**: pre-arms `request.pending_latent` as soon as
+  `<latent>` lands in the output, so the step that PROCESSES the anchor
+  is the fork's entry row (its sample is dropped).
+- **Exit-anchor injection** (`inject_token`): appends `</latent>` as a
+  real sequence token by direct request mutation —
+  `request.append_output_token_ids(tok)` + the worker feed buffer
+  (`input_batch.token_ids_cpu[req_idx, num_tokens_no_spec[req_idx]]` +
+  `is_token_ids`) + re-arming `pending_latent` so the fork drops the
+  exit row's own sample (no double-append); the run loop clears the flag
+  before the next schedule.
+
+## Tools (`orchestrator.tools`)
+
+- `python`: bounded, isolated exec (`python -I`, wall-clock timeout,
+  capped output) capturing stdout/stderr; returns a string
+  (`ERROR: ...` on failure). A resource boundary, not a security sandbox.
+- `search`: Tavily API when `TAVILY_API_KEY` is set; a clear
+  not-configured stub otherwise (no network without the key).
+
+Both follow the reference contract: `(tag, content) -> result text`.
+
+## T7 serve requirements
+
+- env: `VLLM_USE_V2_MODEL_RUNNER=0`, `async_scheduling=False`
+- entry points to install: `vllm.general_plugins` (the T3 backend,
+  `pip install -e /opt/hybridagent-vllm/backend`) and
+  `vllm.logits_processors` (the T4 state machine,
+  `pip install -e /opt/hybridagent-vllm/processor`)
+- this package (no entry points of its own) is imported by the serving
+  process to drive the engine.
+
+Validated by `poc/test_orchestrator_synthetic.py` (tiny HybridAgent model
+through the real fork + backend + state machine: latent cycle, one
+`</python>` tool call, `<result>` re-feed, DONE with coherent text).

serving/orchestrator/__init__.py

@@ -0,0 +1,41 @@
+"""HybridAgent external orchestrator (T5).
+
+Drives a forked in-process vLLM ``LLMEngine`` step-by-step and runs the
+HybridAgent tool barrier on top of the T4 state machine
+(``HybridAgentLogitsProcessor``):
+
+    orch = HybridAgentOrchestrator(
+        model,                       # HybridAgent checkpoint dir
+        tool_executor=...,           # (tag, content) -> result text
+        max_turns=30,
+        **llm_kwargs,                # dtype, enforce_eager, logits_processors, ...
+    )
+    rid = orch.add_request(prompt)   # <latent>-capable sampling params
+    results = orch.run()             # engine.step() cadence + tool barrier
+
+Engine semantics are fixed for the fork: ``VLLM_USE_V2_MODEL_RUNNER=0``
+and ``async_scheduling=False`` (see ``orchestrator.orchestrator``).
+
+Install into the serving venv (no entry points; imported as a package):
+
+    pip install -e /opt/hybridagent-vllm/orchestrator
+"""
+from orchestrator.orchestrator import (
+    BarrierEvent,
+    HybridAgentOrchestrator,
+)
+from orchestrator.tools import (
+    ToolExecutor,
+    make_default_tool_executor,
+    run_python,
+    run_search,
+)
+
+__all__ = [
+    "HybridAgentOrchestrator",
+    "BarrierEvent",
+    "ToolExecutor",
+    "make_default_tool_executor",
+    "run_python",
+    "run_search",
+]

serving/orchestrator/orchestrator.py

@@ -0,0 +1,350 @@
+"""External orchestrator for the forked in-process vLLM engine (T5).
+
+Drives the HybridAgent agentic loop around a real in-process vLLM
+``LLMEngine``:
+
+* builds the engine with the fork's required semantics --
+  ``VLLM_USE_V2_MODEL_RUNNER=0`` (the T2 fork modifies the V1 runner;
+  vLLM 0.30 defaults to V2, where the latent path is silently inactive)
+  and ``async_scheduling=False`` (the async scheduler reserves output
+  placeholders every step, which keeps ``num_new_tokens > 0`` and
+  prevents the fork's zero-token latent gate; L8) -- with an in-process
+  engine core so the T4 state machine's registered handles reach
+  ``scheduler.requests`` and ``model_runner.input_batch.req_ids``;
+* arms the T4 ``HybridAgentLogitsProcessor`` via ``register_engine``
+  with the engine-side request lookup, the worker batch order, the
+  tokenizer, a ``tool_executor`` and the in-process ``inject_token``
+  sequence append;
+* runs the step loop: ``engine.step()`` -> service every request paused
+  in IDLE at the tool barrier (run the tool, tokenize
+  ``<result>...</result>`` with ``add_special_tokens=False``,
+  ``queue_observation``) -> repeat until every request is done.
+
+The two loop behaviors the state machine cannot do from inside the logits
+processor are done here:
+
+1. **Entry-step alignment.** The state machine arms
+   ``request.pending_latent`` one step AFTER it sees ``<latent>``
+   sampled; the fork's entry row is the step that PROCESSES ``<latent>``
+   (its sample is dropped so the position becomes a latent KV position).
+   The orchestrator pre-arms ``pending_latent`` as soon as the entry
+   anchor lands in the output -- the same alignment the T2b e2e harness
+   used.
+
+2. **Exit-anchor injection.** ``inject_token`` appends ``</latent>`` as
+   a real sequence token at a latent exit (direct request mutation, see
+   the method) and re-arms ``pending_latent`` so the fork's exit-row
+   check drops that row's own sample (no double-append of the anchor).
+   The run loop clears ``pending_latent`` before the next schedule, so
+   the following row is a plain token row that feeds ``</latent>``.
+
+The orchestrator also finishes a request at the state machine's DONE
+(e.g. max_turns, fabricated ``<result>``) by aborting it in the engine,
+mirroring the reference engine's stop at DONE.
+"""
+from __future__ import annotations
+
+import os
+import uuid
+from dataclasses import dataclass
+
+# Fork semantics (L6/L8): force the V1 (forked) runner and the in-process
+# engine core. setdefault: an explicit environment wins.
+os.environ.setdefault("VLLM_USE_V2_MODEL_RUNNER", "0")
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
+
+from .tools import ToolExecutor, make_default_tool_executor  # noqa: E402
+
+__all__ = [
+    "HybridAgentOrchestrator",
+    "ToolExecutor",
+    "make_default_tool_executor",
+    "BarrierEvent",
+]
+
+
+@dataclass
+class BarrierEvent:
+    """One tool-barrier service: a request paused in IDLE ran a tool."""
+
+    request_id: str
+    tag: str
+    content: str
+    result: str
+    step: int
+
+
+class HybridAgentOrchestrator:
+    """Wraps an in-process vLLM ``LLMEngine`` and runs the tool barrier.
+
+    Args:
+        model: checkpoint path (a HybridAgent dir: Qwen3 weights +
+            ``adapter.pt``; the T3 backend must be registered in this
+            process, e.g. ``backend.register()`` or the installed
+            ``vllm.general_plugins`` entry point).
+        tool_executor: ``(tag, content) -> result text``; defaults to
+            :func:`make_default_tool_executor` (bounded python +
+            Tavily/stub search).
+        max_turns: the reference turn cap (default 30).
+        llm_kwargs: passed through to ``vllm.LLM`` (``dtype``,
+            ``enforce_eager``, ``max_model_len``,
+            ``gpu_memory_utilization``, ``logits_processors``, ...).
+            ``async_scheduling`` is forced to ``False``.
+    """
+
+    def __init__(
+        self,
+        model: str,
+        *,
+        tool_executor: ToolExecutor | None = None,
+        max_turns: int = 30,
+        **llm_kwargs,
+    ) -> None:
+        from vllm import LLM
+
+        llm_kwargs.pop("async_scheduling", None)
+        llm_kwargs.setdefault("disable_log_stats", True)
+        self.llm = LLM(model=model, async_scheduling=False, **llm_kwargs)
+        self.engine = self.llm.llm_engine
+        # In-process engine core (VLLM_ENABLE_V1_MULTIPROCESSING=0):
+        # engine_core is the InprocClient; its .engine_core is the
+        # EngineCore holding the scheduler and the model executor.
+        self.core = getattr(
+            self.engine.engine_core, "engine_core", self.engine.engine_core
+        )
+        self.scheduler = self.core.scheduler
+        self.runner = self.core.model_executor.driver_worker.model_runner
+        self.tokenizer = self.engine.get_tokenizer()
+        self.max_turns = int(max_turns)
+        self.tool_executor = (
+            tool_executor if tool_executor is not None
+            else make_default_tool_executor()
+        )
+
+        # Fork handle: the exit-row sample keep/drop reads
+        # request.pending_latent through this lookup (None in
+        # multi-process engines, where every latent-row sample is
+        # dropped and the exit anchor is forced in the logits).
+        self.runner._fork_get_request = self.scheduler.requests.get
+
+        from processor import processor as proc_mod
+
+        self._proc = proc_mod
+        self.entry_id = proc_mod.special_token_id(
+            self.tokenizer, proc_mod.ENTRY_ANCHOR
+        )
+        self.exit_id = proc_mod.special_token_id(
+            self.tokenizer, proc_mod.EXIT_ANCHOR
+        )
+        proc_mod.register_engine(
+            get_request=self.scheduler.requests.get,
+            get_req_ids=lambda: list(self.runner.input_batch.req_ids),
+            tokenizer=self.tokenizer,
+            tool_executor=self.tool_executor,
+            inject_token=self.inject_token,
+            max_turns=self.max_turns,
+        )
+
+        self.barrier_log: list[BarrierEvent] = []
+        self.injected: list[tuple[str, int]] = []
+        self._pending_clear: set[str] = set()
+        self._final: dict[str, object] = {}
+        # External request id -> engine-internal id. vLLM 0.30 randomizes
+        # the internal id (assign_request_id); scheduler.requests, the
+        # worker batch and the state machine are keyed by the internal id,
+        # while RequestOutput.request_id carries the external one.
+        self._rid_map: dict[str, str] = {}
+
+    # ------------------------------------------------------------------
+    # requests
+    # ------------------------------------------------------------------
+    def add_request(
+        self,
+        prompt,
+        request_id: str | None = None,
+        sampling_params=None,
+        **params_kwargs,
+    ) -> str:
+        """Submit a prompt to the engine; returns the request id.
+
+        ``sampling_params=None`` builds the ``<latent>``-capable
+        defaults: greedy decoding (``temperature=0``), no
+        penalties/bad-words/structured output (the state machine owns
+        token control), and a ``max_tokens`` budget (``params_kwargs``,
+        default 2048).
+        """
+        from vllm import SamplingParams
+
+        if sampling_params is None:
+            sampling_params = SamplingParams(
+                temperature=0.0,
+                max_tokens=params_kwargs.pop("max_tokens", 2048),
+                **params_kwargs,
+            )
+        ext_rid = request_id or f"ha-{uuid.uuid4().hex[:12]}"
+        internal_rid = self.engine.add_request(ext_rid, prompt, sampling_params)
+        self._rid_map[ext_rid] = internal_rid
+        return ext_rid
+
+    def engine_request_id(self, rid: str) -> str:
+        """The engine-internal request id for an external request id
+        (the key used by ``scheduler.requests``, the worker batch and the
+        state machine)."""
+        return self._rid_map[rid]
+
+    # ------------------------------------------------------------------
+    # in-process sequence append (the T4-registered inject_token)
+    # ------------------------------------------------------------------
+    def inject_token(self, rid: str, tok: int) -> None:
+        """Append ``tok`` to a running sequence (direct request mutation).
+
+        Exact path (single worker, last rank, synchronous scheduling):
+
+        1. engine-side sequence: ``request.append_output_token_ids(tok)``
+           -- ``_output_token_ids``/``_all_token_ids`` grow, so the
+           scheduler counts the token as pending on the next pass and
+           the block hashes extend;
+        2. worker-side feed buffer: ``input_batch.token_ids_cpu[req_idx,
+           pos] = tok`` where ``pos = input_batch.num_tokens_no_spec[
+           req_idx]`` -- the next slot the sampler would have written
+           this step's sample to -- then advance ``num_tokens_no_spec``
+           and mark ``is_token_ids[req_idx, pos] = True`` (a real token
+           slot, not a prompt-embed slot);
+        3. worker-side cache: ``runner.requests[rid].output_token_ids``
+           (penalty bookkeeping; no penalties here, but kept in sync);
+        4. fork sample gate: ``request.pending_latent = True`` so the
+           fork's exit-row check (``keep = not req.pending_latent``)
+           drops this step's own sample -- no double-append of the exit
+           anchor;
+        5. the run loop clears ``pending_latent`` before the next
+           schedule, otherwise the following row would be read as a
+           latent ENTRY row (``pending_latent and num_new_tokens > 0``).
+        """
+        req = self.scheduler.requests[rid]
+        req.append_output_token_ids(tok)
+        ib = self.runner.input_batch
+        idx = ib.req_id_to_index[rid]
+        pos = int(ib.num_tokens_no_spec[idx])
+        ib.token_ids_cpu[idx, pos] = tok
+        ib.is_token_ids[idx, pos] = True
+        ib.num_tokens_no_spec[idx] = pos + 1
+        cached = self.runner.requests.get(rid)
+        if cached is not None:
+            cached.output_token_ids.append(tok)
+        req.pending_latent = True
+        self._pending_clear.add(rid)
+        self.injected.append((rid, int(tok)))
+
+    # ------------------------------------------------------------------
+    # the run loop
+    # ------------------------------------------------------------------
+    def run(self, max_steps: int = 10_000) -> dict:
+        """Drive all submitted requests to completion.
+
+        Returns ``{request_id: RequestOutput}`` for the finished
+        requests.
+        """
+        step_no = 0
+        while self.engine.has_unfinished_requests():
+            if step_no >= max_steps:
+                raise RuntimeError(
+                    f"orchestrator step cap ({max_steps}) reached with "
+                    f"{self.engine.get_num_unfinished_requests()} "
+                    "request(s) unfinished"
+                )
+            self._pre_step()
+            step_no += 1
+            outputs = self.engine.step()
+            self._post_step(outputs, step_no)
+        return dict(self._final)
+
+    def final_text(self, rid: str) -> str:
+        """The engine-side final text of a finished request."""
+        out = self._final.get(rid)
+        if out is None or not getattr(out, "outputs", None):
+            return ""
+        return out.outputs[0].text
+
+    def shutdown(self) -> None:
+        """Shut the in-process engine core down (frees the GPU)."""
+        self.engine.engine_core.shutdown()
+
+    # -- per-step hooks -------------------------------------------------
+    def _pre_step(self) -> None:
+        # Clear the exit-row re-arm (set by inject_token) so the next
+        # schedule reads a plain token row, not a latent entry row.
+        for rid in list(self._pending_clear):
+            req = self.scheduler.requests.get(rid)
+            if req is not None:
+                req.pending_latent = False
+            self._pending_clear.discard(rid)
+
+    def _post_step(self, outputs, step_no: int) -> None:
+        for out in outputs:
+            rid = out.request_id  # external id
+            if out.finished:
+                self._final[rid] = out
+            internal = self._rid_map.get(rid, rid)
+            self._service_barrier(rid, internal, step_no)
+            self._arm_entry(internal)
+            self._finish_if_done(rid, internal)
+
+    def _service_barrier(self, rid: str, internal: str, step_no: int) -> None:
+        """Tool barrier: run the pending tool and re-inject the result."""
+        st = self._proc.get_state(internal)
+        if st is None or st.mode != "IDLE" or st.pending_tool is None:
+            return
+        tag, content = st.pending_tool
+        result = self.tool_executor(tag, content)
+        self.barrier_log.append(
+            BarrierEvent(
+                request_id=rid,
+                tag=tag,
+                content=content,
+                result=result,
+                step=step_no,
+            )
+        )
+        self._proc.record_tool_call(internal, tag)
+        obs_ids = self.tokenizer(
+            f"<result>{result}</result>", add_special_tokens=False
+        )["input_ids"]
+        self._proc.queue_observation(internal, list(obs_ids))
+
+    def _arm_entry(self, rid: str) -> None:
+        """Pre-arm pending_latent once <latent> lands in the output.
+
+        The state machine's own write happens one step later (it detects
+        the entry from the sampled token on the next apply); the fork's
+        entry row is the step that PROCESSES the anchor, so the flag
+        must be set before that row is scheduled.
+        """
+        st = self._proc.get_state(rid)
+        req = self.scheduler.requests.get(rid)
+        if st is None or req is None:
+            return
+        if (
+            st.mode == "DECODE"
+            and req._output_token_ids
+            and req._output_token_ids[-1] == self.entry_id
+            and not req.pending_latent
+        ):
+            req.pending_latent = True
+
+    def _finish_if_done(self, rid: str, internal: str) -> None:
+        """Abort the engine request once the state machine is DONE.
+
+        The reference engine stops at DONE (max_turns, fabricated
+        <result>, budget); vLLM would otherwise keep sampling until its
+        own eos/max_tokens.
+        """
+        st = self._proc.get_state(internal)
+        if st is None or st.mode != "DONE":
+            return
+        req = self.scheduler.requests.get(internal)
+        if req is None:
+            return
+        from vllm.v1.request import RequestStatus
+
+        if req.status == RequestStatus.RUNNING:
+            self.engine.abort_request([rid])

serving/orchestrator/pyproject.toml

@@ -0,0 +1,15 @@
+[build-system]
+requires = ["setuptools>=68"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hybridagent-orchestrator"
+version = "0.1.0"
+description = "External orchestrator: drives the forked in-process vLLM engine and runs the HybridAgent tool barrier (T5)"
+requires-python = ">=3.10"
+# No entry points: this package is imported by the serving process. The
+# engine-side plugins are the T3 backend (vllm.general_plugins) and the
+# T4 state machine (vllm.logits_processors).
+
+[tool.setuptools]
+packages = ["orchestrator"]

serving/orchestrator/tools.py

@@ -0,0 +1,134 @@
+"""Tool executors for the HybridAgent tool barrier (T5).
+
+Reference parity (togethercomputer/HybridAgent ``ToolExecutor``): one tool
+call is ``(tag, content) -> result text``. The orchestrator wraps the
+returned text in ``<result>...</result>``, tokenizes it and re-injects it
+as FEED tokens for the paused request.
+
+Two tools (``tag`` in {"python", "search"}, the state machine's
+``TOOL_TAGS``):
+
+* ``python`` -- runs ``content`` as Python in a bounded, isolated
+  interpreter (``python -I``: no PYTHONPATH, no environment, no user
+  site-packages), with a wall-clock timeout and a capped output. Returns
+  the captured stdout, or an ``ERROR: ...`` string on failure/timeout.
+  This is a resource boundary (time, output, environment isolation), not
+  a security sandbox: the code still runs with the process's filesystem
+  permissions.
+* ``search`` -- web search via the Tavily API when ``TAVILY_API_KEY`` is
+  set in the environment; otherwise a clear not-configured stub string.
+  No network activity happens without the key.
+
+Both return plain strings (the reference contract).
+"""
+from __future__ import annotations
+
+import json
+import os
+import subprocess
+import sys
+import urllib.error
+import urllib.request
+
+__all__ = [
+    "ToolExecutor",
+    "make_default_tool_executor",
+    "run_python",
+    "run_search",
+]
+
+PYTHON_TIMEOUT_S = 10.0
+PYTHON_MAX_OUTPUT = 8192
+SEARCH_TIMEOUT_S = 20.0
+SEARCH_MAX_OUTPUT = 8192
+TAVILY_ENDPOINT = "https://api.tavily.com/search"
+
+
+def _truncate(text: str, limit: int) -> str:
+    return text if len(text) <= limit else text[:limit] + "...<truncated>"
+
+
+def run_python(
+    code: str,
+    *,
+    timeout: float = PYTHON_TIMEOUT_S,
+    max_output: int = PYTHON_MAX_OUTPUT,
+) -> str:
+    """Run ``code`` in a bounded, isolated interpreter; return a string."""
+    if not isinstance(code, str) or not code.strip():
+        return "ERROR: empty python code"
+    env = {"PATH": os.environ.get("PATH", "/usr/bin:/bin")}
+    try:
+        proc = subprocess.run(
+            [sys.executable, "-I", "-c", code],
+            capture_output=True,
+            text=True,
+            timeout=timeout,
+            env=env,
+        )
+    except subprocess.TimeoutExpired:
+        return f"ERROR: python tool timed out after {timeout:.0f}s"
+    out = (proc.stdout or "").strip()
+    if proc.returncode != 0:
+        err_lines = (proc.stderr or "").strip().splitlines()
+        last = err_lines[-1] if err_lines else f"exit code {proc.returncode}"
+        return f"ERROR: {last}"
+    return _truncate(out, max_output)
+
+
+def run_search(
+    query: str,
+    *,
+    timeout: float = SEARCH_TIMEOUT_S,
+    max_output: int = SEARCH_MAX_OUTPUT,
+) -> str:
+    """Web search via Tavily when configured; a stub string otherwise."""
+    key = os.environ.get("TAVILY_API_KEY", "").strip()
+    if not key:
+        return (
+            "web_search not configured: TAVILY_API_KEY is not set; no results."
+        )
+    if not isinstance(query, str) or not query.strip():
+        return "ERROR: empty search query"
+    body = json.dumps(
+        {"api_key": key, "query": query, "max_results": 3}
+    ).encode("utf-8")
+    req = urllib.request.Request(
+        TAVILY_ENDPOINT,
+        data=body,
+        headers={"Content-Type": "application/json"},
+    )
+    try:
+        with urllib.request.urlopen(req, timeout=timeout) as resp:
+            data = json.loads(resp.read().decode("utf-8"))
+    except (urllib.error.URLError, OSError, ValueError) as exc:
+        return f"ERROR: search failed: {exc}"
+    parts: list[str] = []
+    if data.get("answer"):
+        parts.append(str(data["answer"]))
+    for result in (data.get("results") or [])[:3]:
+        title = result.get("title") or ""
+        content = result.get("content") or ""
+        parts.append(f"{title}: {content}" if title else content)
+    out = "\n".join(p for p in parts if p).strip()
+    return _truncate(out, max_output) if out else "no results"
+
+
+class ToolExecutor:
+    """``(tag, content) -> result text`` for the orchestrator's tool barrier."""
+
+    def __init__(self, python=None, search=None) -> None:
+        self.python = python if python is not None else run_python
+        self.search = search if search is not None else run_search
+
+    def __call__(self, tag: str, content: str) -> str:
+        if tag == "python":
+            return self.python(content)
+        if tag == "search":
+            return self.search(content)
+        return f"ERROR: unknown tool {tag!r}"
+
+
+def make_default_tool_executor() -> ToolExecutor:
+    """The default barrier executor (bounded python + Tavily/stub search)."""
+    return ToolExecutor()

serving/processor/__init__.py

@@ -0,0 +1,62 @@
+"""HybridAgent vLLM stateful logits processor (T4).
+
+The DECODE/LATENT/FEED/IDLE/DONE generation state machine from
+togethercomputer/HybridAgent running inside vLLM as a stateful logits
+processor (``vllm.logits_processors`` entry-point group). Drives the
+fork's zero-token latent steps via ``request.pending_latent`` and
+forces tokens (latent exit anchor, FEED observations) via logits
+masking.
+
+Install into the serving venv to activate:
+
+    pip install -e /opt/hybridagent-vllm/processor
+
+T5 (the external orchestrator) arms it once, in the engine process:
+
+    import processor
+    processor.register_engine(
+        get_request=...,        # engine-side Request by req_id
+        get_req_ids=...,        # current worker batch order
+        tokenizer=...,          # <latent>/</latent> resolved at runtime
+        tool_executor=...,      # (tag, content) -> result text
+        inject_token=...,       # in-process sequence append (exit anchor)
+        max_turns=30,
+    )
+    # at each tool barrier (request paused in IDLE):
+    tag, content = processor.get_state(req_id).pending_tool
+    result = tool_executor(tag, content)
+    processor.queue_observation(
+        req_id,
+        tokenizer(f"<result>{result}</result>",
+                  add_special_tokens=False)["input_ids"],
+    )
+"""
+from processor.processor import (
+    HybridAgentLogitsProcessor,
+    get_state,
+    queue_observation,
+    register_engine,
+    record_tool_call,
+)
+from processor.state_machine import (
+    DECODE,
+    DONE,
+    FEED,
+    IDLE,
+    LATENT,
+    RequestState,
+)
+
+__all__ = [
+    "HybridAgentLogitsProcessor",
+    "register_engine",
+    "queue_observation",
+    "record_tool_call",
+    "get_state",
+    "RequestState",
+    "DECODE",
+    "LATENT",
+    "FEED",
+    "IDLE",
+    "DONE",
+]

serving/processor/processor.py

@@ -0,0 +1,370 @@
+"""vLLM stateful logits processor: the HybridAgent state machine (T4).
+
+Wired through the ``vllm.logits_processors`` entry-point group (see
+pyproject.toml):
+
+    [project.entry-points."vllm.logits_processors"]
+    hybridagent_state_machine = "processor.processor:HybridAgentLogitsProcessor"
+
+The processor mirrors the ``AdapterLogitsProcessor`` pattern
+(``vllm/v1/sample/logits_processor/``): a batch-index -> state map
+maintained from ``BatchUpdate`` (the ``output_tok_ids`` entries are
+live references to each request's running output list), advanced in
+``apply()`` from the live token stream and the current logits' argmax.
+
+Fork integration (branch ``hybridagent-latent-fork``):
+
+* ``request.pending_latent`` -- written on the engine-side ``Request``
+  (resolved through the T5-registered request store) so the fork's
+  scheduler gate (``scheduler.py``: ``pending_latent and
+  num_new_tokens == 0 -> latent_positions = 1``) schedules the
+  zero-token latent step. vLLM 0.30's ``BatchUpdate`` carries only
+  ``SamplingParams`` + output-token references and its
+  ``ForwardContext`` exposes no model-runner handle, so the engine-side
+  request is reached through the in-process store T5 registers (the T5
+  deployment is the in-process ``LLMEngine``).
+* Forced tokens -- the logits row is masked to a single id (latent
+  exit anchor fallback, FEED observation tokens).
+* Exit-anchor injection -- in-process engines: T5-registered
+  ``inject_token(req_id, token_id)`` appends ``</latent>`` to the
+  sequence at a latent exit (reference parity: ``<latent>``, n latent
+  forwards, ``</latent>``).
+
+T4b -- reference ``blocked_token_sets`` (engine.py:140-149): every
+DECODE (text) logits row is masked to ``-inf`` at the blocked anchor
+ids, from the first step: in ``<think>`` block ``</latent>`` (and
+``<think>``), out of ``<think>`` block ``</latent>``, ``<latent>``
+(and ``</think>``). ``in_think`` is initialized from the prompt
+(engine.py:221) and tracked per request from the generated
+``<think>``/``</think>`` tokens; LATENT exit readouts stay raw. The
+masks are applied before the single-token forcing, so FEED/exit
+forcing is unchanged. Non-agent requests still get no
+``pending_latent`` writes.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import TYPE_CHECKING, Any, Callable
+
+import torch
+
+from vllm import SamplingParams
+from vllm.v1.sample.logits_processor import (
+    BatchUpdate,
+    LogitsProcessor,
+    process_dict_updates,
+)
+
+from processor.state_machine import (
+    DEFAULT_MAX_TURNS,
+    ENTRY_ANCHOR,
+    EXIT_ANCHOR,
+    THINK_CLOSE,
+    THINK_OPEN,
+    TOOL_TAG_TAIL_TOKENS,
+    RequestState,
+    StepAction,
+    advance,
+    queue_observation as _sm_queue_observation,
+    record_tool_call as _sm_record_tool_call,
+)
+
+if TYPE_CHECKING:
+    from vllm.config import VllmConfig
+
+__all__ = [
+    "HybridAgentLogitsProcessor",
+    "register_engine",
+    "queue_observation",
+    "record_tool_call",
+    "get_state",
+    "special_token_id",
+]
+
+# Live processor instances in this process (one worker, one instance).
+_PROCESSORS: list["HybridAgentLogitsProcessor"] = []
+
+
+def special_token_id(tokenizer: Any, token: str) -> int:
+    """Id of a special token; fails loudly if the tokenizer does not know
+    it (reference ``constants.special_token_id``)."""
+    tok_id = tokenizer.convert_tokens_to_ids(token)
+    unk_id = getattr(tokenizer, "unk_token_id", None)
+    if tok_id is None or (unk_id is not None and tok_id == unk_id):
+        raise ValueError(
+            f"Tokenizer does not know the special token {token!r}; "
+            "load a HybridAgent checkpoint or add the anchor tokens first."
+        )
+    return int(tok_id)
+
+
+def _optional_token_id(tokenizer: Any, token: str) -> int | None:
+    """T4b: token id or ``None`` when the tokenizer does not know it.
+
+    The think-tag guards (reference ``think_guards``) degrade to
+    anchor-only masks in that case (reference ``think_guards=False``
+    path); they never block arming the engine.
+    """
+    try:
+        tok_id = tokenizer.convert_tokens_to_ids(token)
+    except Exception:
+        return None
+    unk_id = getattr(tokenizer, "unk_token_id", None)
+    if tok_id is None or (unk_id is not None and tok_id == unk_id):
+        return None
+    return int(tok_id)
+
+
+@dataclass
+class EngineHandles:
+    """Process-level handles T5 registers once (in-process engine)."""
+
+    get_request: Callable[[str], Any] | None = None
+    get_req_ids: Callable[[], list[str] | None] | None = None
+    tokenizer: Any | None = None
+    tool_executor: Callable[[str, str], str] | None = None
+    inject_token: Callable[[str, int], None] | None = None
+    max_turns: int = DEFAULT_MAX_TURNS
+    entry_id: int | None = None
+    exit_id: int | None = None
+    think_open_id: int | None = None
+    think_close_id: int | None = None
+    eos_token_ids: frozenset[int] = field(default_factory=frozenset)
+
+
+_engine: EngineHandles | None = None
+
+
+def register_engine(
+    *,
+    get_request: Callable[[str], Any] | None = None,
+    get_req_ids: Callable[[], list[str] | None] | None = None,
+    tokenizer: Any | None = None,
+    tool_executor: Callable[[str, str], str] | None = None,
+    inject_token: Callable[[str, int], None] | None = None,
+    max_turns: int = DEFAULT_MAX_TURNS,
+    latent_open_id: int | None = None,
+    latent_close_id: int | None = None,
+) -> None:
+    """T5 interface: arm the state machine in this process.
+
+    Args:
+        get_request: engine-side ``Request`` by req_id (the object the
+            fork's scheduler reads ``pending_latent`` from).
+        get_req_ids: current worker batch order (logits row i is
+            ``req_ids[i]``).
+        tokenizer: the HybridAgent tokenizer; the ``<latent>`` /
+            ``</latent>`` ids are resolved from it at runtime (explicit
+            ``latent_open_id`` / ``latent_close_id`` override). The
+            T4b think-tag ids (``<think>`` / ``</think>``) are resolved
+            leniently: unknown tag -> anchor-only masks.
+        tool_executor: ``(tag, content) -> result text`` for one
+            pending tool call (tag in {"python", "search"}).
+        inject_token: in-process sequence append used to feed the exit
+            anchor at a latent exit (reference parity).
+        max_turns: the reference turn cap (default 30).
+    """
+    global _engine
+    entry_id = latent_open_id
+    exit_id = latent_close_id
+    think_open: int | None = None
+    think_close: int | None = None
+    eos: set[int] = set()
+    if tokenizer is not None:
+        if entry_id is None:
+            entry_id = special_token_id(tokenizer, ENTRY_ANCHOR)
+        if exit_id is None:
+            exit_id = special_token_id(tokenizer, EXIT_ANCHOR)
+        # T4b: think-tag ids for the blocked-set guards (lenient:
+        # unknown tag -> anchor-only masks).
+        think_open = _optional_token_id(tokenizer, THINK_OPEN)
+        think_close = _optional_token_id(tokenizer, THINK_CLOSE)
+        cand = getattr(tokenizer, "eos_token_id", None)
+        if cand is not None:
+            eos.update(cand if isinstance(cand, (list, tuple)) else [cand])
+    if entry_id is None or exit_id is None:
+        raise ValueError(
+            "register_engine: the <latent>/</latent> ids must be "
+            "resolvable from the tokenizer (or passed explicitly)."
+        )
+    _engine = EngineHandles(
+        get_request=get_request,
+        get_req_ids=get_req_ids,
+        tokenizer=tokenizer,
+        tool_executor=tool_executor,
+        inject_token=inject_token,
+        max_turns=max_turns,
+        entry_id=int(entry_id),
+        exit_id=int(exit_id),
+        think_open_id=think_open,
+        think_close_id=think_close,
+        eos_token_ids=frozenset(int(i) for i in eos),
+    )
+
+
+def _find_state(req_id: str) -> RequestState | None:
+    eng = _engine
+    if eng is None or eng.get_req_ids is None:
+        return None
+    req_ids = eng.get_req_ids()
+    if req_ids is None or req_id not in req_ids:
+        return None
+    idx = req_ids.index(req_id)
+    for proc in _PROCESSORS:
+        if idx in proc.info:
+            return proc.info[idx][0]
+    return None
+
+
+def queue_observation(req_id: str, token_ids: list[int]) -> None:
+    """T5 tool barrier: feed a ``<result>...`` observation to a request
+    paused in IDLE (empty list = empty observation)."""
+    state = _find_state(req_id)
+    if state is None:
+        raise KeyError(f"no tracked HybridAgent state for request {req_id!r}")
+    _sm_queue_observation(state, token_ids)
+
+
+def record_tool_call(req_id: str, tag: str) -> None:
+    """T5: report an executed tool call (per-tag counter bookkeeping)."""
+    state = _find_state(req_id)
+    if state is None:
+        raise KeyError(f"no tracked HybridAgent state for request {req_id!r}")
+    _sm_record_tool_call(state, tag)
+
+
+def get_state(req_id: str) -> RequestState | None:
+    """T5: read a request's state (mode, pending_tool, end_reason,
+    tool counters)."""
+    return _find_state(req_id)
+
+
+class HybridAgentLogitsProcessor(LogitsProcessor):
+    """Stateful logits processor running the HybridAgent state machine."""
+
+    def __init__(
+        self,
+        vllm_config: "VllmConfig",
+        device: torch.device,
+        is_pin_memory: bool,
+    ):
+        self.device = device
+        # batch index -> (state, live output-token list, max_tokens,
+        # stop ids)
+        self.info: dict[
+            int, tuple[RequestState, list[int], int | None, frozenset[int]]
+        ] = {}
+        _PROCESSORS.append(self)
+
+    def is_argmax_invariant(self) -> bool:
+        # Forced tokens (latent exit, FEED) change greedy argmax.
+        return False
+
+    @staticmethod
+    def _new_info(
+        params: SamplingParams,
+        prompt_ids: list[int] | None,
+        output_ids: list[int],
+    ) -> tuple[RequestState, list[int], int | None, frozenset[int]]:
+        max_turns = _engine.max_turns if _engine is not None else DEFAULT_MAX_TURNS
+        state = RequestState(max_turns=max_turns)
+        # T4b: reference in_think initialization (engine.py:221/239) --
+        # the chat template ends in the <think> token.
+        eng = _engine
+        if eng is not None and eng.tokenizer is not None and prompt_ids:
+            prompt_text = eng.tokenizer.decode(
+                list(prompt_ids), skip_special_tokens=False
+            )
+            state.in_think = prompt_text.rstrip("\n").endswith(THINK_OPEN)
+        return (
+            state,
+            output_ids,
+            params.max_tokens,
+            frozenset(params.all_stop_token_ids),
+        )
+
+    def update_state(self, batch_update: BatchUpdate | None):
+        process_dict_updates(self.info, batch_update, self._new_info)
+
+    def apply(self, logits: torch.Tensor) -> torch.Tensor:
+        eng = _engine
+        if not self.info or eng is None or eng.get_req_ids is None:
+            return logits
+        req_ids = eng.get_req_ids()
+        if req_ids is None:
+            return logits
+        tracked = [idx for idx in self.info if idx < len(req_ids)]
+        if not tracked:
+            return logits
+        argmax_ids = torch.argmax(logits[tracked], dim=-1).tolist()
+
+        forced: list[tuple[int, int]] = []
+        masked: list[tuple[int, frozenset[int]]] = []
+        for idx, argmax_id in zip(tracked, argmax_ids):
+            st, outputs, max_tokens, stop_ids = self.info[idx]
+            last_token = outputs[-1] if outputs else None
+            tail = ""
+            full_text_getter = None
+            if st.mode == "DECODE" and eng.tokenizer is not None:
+                # Reference order: append the just-sampled token to the
+                # text-since-resume window, then take the tail of the last
+                # ~50 *text* tokens (observations/anchors excluded).
+                window = list(st.text_ids_since_resume)
+                if last_token is not None and last_token != eng.exit_id:
+                    window.append(last_token)
+                if window:
+                    tail = eng.tokenizer.decode(
+                        window[-TOOL_TAG_TAIL_TOKENS:],
+                        skip_special_tokens=False,
+                    )
+                    full_text_getter = eng.tokenizer.decode
+
+            action = advance(
+                st,
+                last_token=last_token,
+                argmax_id=argmax_id,
+                tail_text=tail,
+                full_text_getter=(
+                    (lambda: full_text_getter(outputs,
+                                                skip_special_tokens=False))
+                    if tail
+                    else None
+                ),
+                entry_id=eng.entry_id,
+                exit_id=eng.exit_id,
+                stop_ids=stop_ids | eng.eos_token_ids,
+                max_tokens=max_tokens,
+                output_len=len(outputs),
+                think_open_id=eng.think_open_id,
+                think_close_id=eng.think_close_id,
+            )
+
+            req = (
+                eng.get_request(req_ids[idx])
+                if eng.get_request is not None
+                else None
+            )
+            if req is not None and st.dirty:
+                req.pending_latent = action.pending_latent
+            if action.inject_token is not None and eng.inject_token is not None:
+                eng.inject_token(req_ids[idx], action.inject_token)
+                action = StepAction(pending_latent=action.pending_latent)
+            if action.blocked_ids:
+                masked.append((idx, action.blocked_ids))
+            if action.force_id is not None:
+                forced.append((idx, action.force_id))
+
+        # T4b: reference text-mode anchor masks (engine.py:443-451 +
+        # sampling.py:75 masked_fill): block </latent> in all text
+        # mode, <latent> outside <think> (plus the think-tag guards).
+        # Applied before the forcing loop so forced tokens (FEED,
+        # latent exit) still win on their own rows.
+        if masked:
+            dev = logits.device
+            for idx, ids in masked:
+                logits[idx, torch.tensor(sorted(ids), device=dev,
+                                          dtype=torch.long)] = float("-inf")
+        for idx, tok in forced:
+            logits[idx].fill_(float("-inf"))
+            logits[idx, tok] = 0.0
+        return logits

serving/processor/pyproject.toml

@@ -0,0 +1,24 @@
+[build-system]
+requires = ["setuptools>=61"]
+build-backend = "setuptools.build_meta"
+
+[project]
+name = "hybridagent-vllm-processor"
+version = "0.1.0"
+description = "vLLM stateful logits processor for HybridAgent: DECODE/LATENT/FEED/IDLE/DONE state machine + forced </latent> (zero-token latent steps)"
+requires-python = ">=3.10"
+# vllm and torch are intentionally not declared: they come from the
+# serving environment, mirroring hybridagent-vllm-backend.
+dependencies = []
+
+[project.entry-points."vllm.logits_processors"]
+hybridagent_state_machine = "processor.processor:HybridAgentLogitsProcessor"
+
+[tool.setuptools]
+# Flat layout: the `processor` package files (__init__.py, processor.py,
+# state_machine.py) live in the project root alongside pyproject.toml, not in
+# a nested processor/processor/ dir. Map the package to "." so the
+# vllm.logits_processors entry point (processor.processor:HybridAgentLogitsProcessor)
+# resolves after an editable install.
+package-dir = {"processor" = "."}
+packages = ["processor"]

serving/processor/state_machine.py

@@ -0,0 +1,410 @@
+"""HybridAgent generation state machine (T4) -- pure control flow.
+
+Mirrors the control flow of togethercomputer/HybridAgent's
+``HybridGenerationEngine`` (``hybrid_agent/engine.py`` @ HEAD 1b41cebb):
+the DECODE/LATENT/FEED/IDLE/DONE mode machine, latent entry/exit
+anchors, the tool barrier, observation re-injection, and the turn cap.
+
+T4b -- reference ``blocked_token_sets`` port (engine.py:140-149):
+every DECODE step now carries the text-mode anchor masks
+(``StepAction.blocked_ids``): in ``<think>`` block
+``{</latent>, <think>}``, out of ``<think>`` block
+``{</latent>, <latent>, </think>}`` (guards active when the think-tag
+ids are known, as in the reference default). ``in_think`` is tracked
+per request from the ``<think>``/``</think>`` tokens in the generated
+sequence (reference engine.py:493-496), initialized from the prompt
+(engine.py:221/239) and reset when an observation is fully fed
+(engine.py:434/590). LATENT exit readouts stay raw (engine.py:373),
+so ``</latent>`` can only be produced by the forced latent-exit
+token and ``<latent>`` only inside ``<think>``.
+
+No vLLM/torch dependency: the vLLM-side glue (logits masking,
+``request.pending_latent`` writes, entry-point wiring) lives in
+``processor.processor``. This module is what the unit test
+(``poc/test_state_machine.py``) drives directly.
+
+Per-step contract -- the caller (the logits processor) invokes
+``advance()`` once per generation step, after the forward pass and
+before sampling, with:
+
+* ``last_token``  -- the token sampled at the previous step (``None``
+  before the first sample; unchanged across zero-token latent steps,
+  since the fork drops latent-row samples);
+* ``argmax_id``   -- the argmax of the current step's logits row;
+* ``tail_text``   -- decoded text of the last ~50 text tokens (the
+  reference tool-tag window);
+* ``full_text_getter`` -- lazily decoded full output text, used only
+  when a tool tag fires (content extraction, reference semantics).
+
+Latent-phase timing (reference parity): the exit decision that the
+reference makes at the *start* of a step (previous step's logits,
+``latent_steps >= 4``) is made here at the *end* of the 4th latent
+forward -- the point where the same counter and the same logits are
+first both available in vLLM. The exit anchor is then injected into
+the sequence so the next step feeds ``</latent>`` as a real token:
+``<latent>``, 4 latent forwards, ``</latent>``.
+
+The returned :class:`StepAction` tells the caller what to do with the
+current logits row and with the request's fork fields.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Callable
+
+DECODE, LATENT, FEED, IDLE, DONE = "DECODE", "LATENT", "FEED", "IDLE", "DONE"
+
+# HybridAgent trajectory format (reference hybrid_agent/constants.py).
+ENTRY_ANCHOR = "<latent>"
+EXIT_ANCHOR = "</latent>"
+THINK_OPEN = "<think>"
+THINK_CLOSE = "</think>"
+RESULT_OPEN = "<result>"
+TOOL_TAGS = ("search", "python")
+TOOL_STOP_TAGS = tuple(f"</{tag}>" for tag in TOOL_TAGS)
+
+MIN_LATENT_STEPS = 4
+MAX_LATENT_STEPS = 200
+TOOL_TAG_TAIL_TOKENS = 50
+DEFAULT_MAX_TURNS = 30
+
+
+@dataclass
+class StepAction:
+    """What the caller must do with the current step.
+
+    Attributes:
+        force_id: mask the current logits row to this single token
+            (FEED observation tokens; latent-exit fallback when no
+            sequence injector is available).
+        pending_latent: value to write to ``request.pending_latent``
+            for the next scheduler pass (fork zero-token step gate).
+        inject_token: token id to append to the sequence *now* (the
+            exit anchor at a latent exit; in-process engines only).
+        blocked_ids: T4b -- token ids to mask to ``-inf`` in the
+            current logits row: the reference ``blocked_token_sets``
+            text-mode masks (``engine.py:140-149``). Set on DECODE
+            steps only; ``None`` elsewhere (LATENT exit readouts stay
+            raw, FEED/IDLE/DONE rows are forced/padded).
+    """
+
+    force_id: int | None = None
+    pending_latent: bool = False
+    inject_token: int | None = None
+    blocked_ids: frozenset[int] | None = None
+
+
+class RequestState:
+    """Per-request state of one HybridAgent trajectory row."""
+
+    def __init__(
+        self,
+        max_turns: int = DEFAULT_MAX_TURNS,
+        min_latent_steps: int = MIN_LATENT_STEPS,
+        max_latent_steps: int = MAX_LATENT_STEPS,
+    ):
+        self.mode: str = DECODE
+        # T4b: inside the <think>...</think> region? (reference
+        # ``RowState.in_think``; initialized from the prompt by the
+        # caller, then tracked from the <think>/</think> tokens in the
+        # generated sequence -- engine.py:221/239/493-496).
+        self.in_think: bool = False
+        self.max_turns = int(max_turns)
+        self.min_latent_steps = int(min_latent_steps)
+        self.max_latent_steps = int(max_latent_steps)
+        # steps of the open latent span (completed latent forwards)
+        self.latent_steps: int = 0
+        self.turn_count: int = 1
+        self.feed_queue: list[int] = []
+        self.pending_tool: tuple[str, str] | None = None
+        self.python_calls: int = 0
+        self.search_calls: int = 0
+        self.tool_requests: int = 0
+        self.anchor_exits: int = 0
+        self.forced_exits: int = 0
+        self.end_reason: str = ""
+        # has left plain DECODE; gates the pending_latent writes so
+        # ordinary requests are a strict no-op
+        self.dirty: bool = False
+        # the <latent> token id captured at entry; a real zero-token
+        # step is one whose last sampled token is unchanged
+        self._entry_token: int | None = None
+        # text tokens since the last resume (tool-tag window source)
+        self.text_ids_since_resume: list[int] = []
+
+    @property
+    def is_done(self) -> bool:
+        return self.mode == DONE
+
+
+def find_tool_tag(tail_text: str) -> str | None:
+    """First closing tool tag in the decoded tail (reference order)."""
+    for stop in TOOL_STOP_TAGS:
+        if stop in tail_text:
+            return stop
+    return None
+
+
+def extract_content(text: str, tag: str) -> str:
+    """Extract text between the *last* ``<tag>...</tag>`` pair
+    (reference ``ToolExecutor.extract_content``)."""
+    start_tag = f"<{tag}>"
+    end_tag = f"</{tag}>"
+    try:
+        end_pos = text.rindex(end_tag)
+        start_pos = text.rindex(start_tag, 0, end_pos)
+        return text[start_pos + len(start_tag) : end_pos].strip()
+    except ValueError:
+        return ""
+
+
+def _finish(state: RequestState, reason: str) -> None:
+    state.mode = DONE
+    state.end_reason = state.end_reason or reason
+    state.dirty = True
+
+
+def _exit_latent(state: RequestState, action: StepAction, exit_id: int,
+                 forced: bool) -> None:
+    state.mode = DECODE
+    state.latent_steps = 0
+    state.dirty = True
+    if forced:
+        state.forced_exits += 1
+    else:
+        state.anchor_exits += 1
+    action.pending_latent = False
+    action.inject_token = exit_id
+    # Fallback for engines without an injector (multi-process): mask the
+    # current row to the exit anchor. The fork's latent-row sample filter
+    # drops that sample, so this path is best-effort; the in-process
+    # inject_token path is the parity-correct one.
+    action.force_id = exit_id
+
+
+def _advance_latent(
+    state: RequestState,
+    action: StepAction,
+    last_token: int | None,
+    argmax_id: int,
+    exit_id: int,
+) -> None:
+    if last_token != state._entry_token:
+        # Not a zero-token step: the fork scheduled a normal token step
+        # (the T2 entry bridge is absent, so the entry step's sample was
+        # appended). Do not count it as a latent forward and do not exit;
+        # keep the loop pending and decode undisturbed.
+        action.pending_latent = True
+        return
+    # A latent forward just completed (the fork drops its sample).
+    state.latent_steps += 1
+    # Reference exit order: forced cap first, then the anchor readout.
+    if (
+        state.max_latent_steps > 0
+        and state.latent_steps >= state.max_latent_steps
+    ):
+        _exit_latent(state, action, exit_id, forced=True)
+        return
+    if (
+        state.latent_steps >= state.min_latent_steps
+        and argmax_id == exit_id
+    ):
+        _exit_latent(state, action, exit_id, forced=False)
+        return
+    action.pending_latent = True
+
+
+def _blocked_token_set(
+    state: RequestState,
+    entry_id: int,
+    exit_id: int,
+    think_open_id: int | None,
+    think_close_id: int | None,
+) -> frozenset[int]:
+    """T4b: the reference ``blocked_token_sets`` (engine.py:140-145)
+    for the current think state, with the think guards active when the
+    think-tag ids are known (the reference default ``think_guards``):
+
+    * in-think  -> block ``</latent>`` (and ``<think>``)
+    * out-think -> block ``</latent>``, ``<latent>`` (and ``</think>``)
+
+    so ``</latent>`` can only come from the forced latent-exit readout
+    and ``<latent>`` only inside ``<think>``.
+    """
+    if state.in_think:
+        ids = {exit_id}
+        if think_open_id is not None:
+            ids.add(think_open_id)
+    else:
+        ids = {exit_id, entry_id}
+        if think_close_id is not None:
+            ids.add(think_close_id)
+    return frozenset(ids)
+
+
+def _advance_feed(state: RequestState, action: StepAction) -> None:
+    # One forced observation token per step (reference FEED).
+    tok = state.feed_queue[0]
+    state.feed_queue.pop(0)
+    action.force_id = tok
+    state.dirty = True
+    if not state.feed_queue:
+        # Observation fully fed: the reference starts the next turn here.
+        state.text_ids_since_resume = []
+        state.in_think = False  # T4b (reference engine.py:434)
+        if state.turn_count > state.max_turns:
+            _finish(state, "max_turns")
+        else:
+            state.mode = DECODE
+
+
+def _advance_decode(
+    state: RequestState,
+    action: StepAction,
+    last_token: int | None,
+    tail_text: str,
+    full_text_getter: Callable[[], str] | None,
+    entry_id: int,
+    exit_id: int,
+    stop_ids: frozenset[int],
+    max_tokens: int | None,
+    output_len: int,
+    think_open_id: int | None = None,
+    think_close_id: int | None = None,
+) -> None:
+    # Budget gate (vLLM enforces max_tokens natively too; this keeps the
+    # mode bookkeeping consistent).
+    if max_tokens is not None and output_len >= max_tokens:
+        _finish(state, "budget")
+        return
+    if last_token is None:
+        # T4b: first text step after prefill -- the reference masks it
+        # from the start (engine.py:443-451) with the prompt-initialized
+        # in_think state.
+        action.blocked_ids = _blocked_token_set(
+            state, entry_id, exit_id, think_open_id, think_close_id
+        )
+        return
+    if last_token == entry_id:
+        # DECODE -> LATENT: <latent> is a real token in the sequence;
+        # the current step is the entry (latent) step, whose readout
+        # stays raw (reference engine.py:373/446).
+        state.mode = LATENT
+        state.latent_steps = 0
+        state._entry_token = entry_id
+        state.dirty = True
+        action.pending_latent = True
+        return
+    if last_token in stop_ids:
+        _finish(state, "eos")
+        return
+    # T4b: think-region tracking from the generated sequence
+    # (reference engine.py:493-496).
+    if think_open_id is not None and last_token == think_open_id:
+        state.in_think = True
+    elif think_close_id is not None and last_token == think_close_id:
+        state.in_think = False
+    # T4b: reference text-mode anchor masks on EVERY text step
+    # (engine.py:443-451): </latent> always, <latent> outside <think> --
+    # so the anchors are only produced by the state machine's forced
+    # tokens.
+    action.blocked_ids = _blocked_token_set(
+        state, entry_id, exit_id, think_open_id, think_close_id
+    )
+    if last_token != exit_id:
+        # The exit anchor is fed, not model text; keep it out of the
+        # tool-tag window.
+        state.text_ids_since_resume.append(last_token)
+    # Fabricated-observation guard (reference reject_fabricated_result).
+    if RESULT_OPEN in tail_text:
+        _finish(state, "invalid:fabricated_result")
+        return
+    stop = find_tool_tag(tail_text)
+    if stop is not None:
+        tag = stop[2:-1]
+        full_text = (
+            full_text_getter() if full_text_getter is not None else tail_text
+        )
+        state.pending_tool = (tag, extract_content(full_text, tag))
+        state.tool_requests += 1
+        state.mode = IDLE
+        state.dirty = True
+
+
+def advance(
+    state: RequestState,
+    *,
+    last_token: int | None,
+    argmax_id: int,
+    tail_text: str = "",
+    full_text_getter: Callable[[], str] | None = None,
+    entry_id: int,
+    exit_id: int,
+    stop_ids: frozenset[int] = frozenset(),
+    max_tokens: int | None = None,
+    output_len: int = 0,
+    think_open_id: int | None = None,
+    think_close_id: int | None = None,
+) -> StepAction:
+    """Advance one request's state machine by one generation step."""
+    action = StepAction()
+    if state.mode == DONE:
+        return action
+    if state.mode == LATENT:
+        _advance_latent(state, action, last_token, argmax_id, exit_id)
+        return action
+    if state.mode == FEED:
+        _advance_feed(state, action)
+        return action
+    if state.mode == IDLE:
+        # Paused at the tool barrier; T5 resumes via queue_observation().
+        return action
+    _advance_decode(
+        state,
+        action,
+        last_token,
+        tail_text,
+        full_text_getter,
+        entry_id,
+        exit_id,
+        stop_ids,
+        max_tokens,
+        output_len,
+        think_open_id,
+        think_close_id,
+    )
+    return action
+
+
+def queue_observation(state: RequestState, token_ids: list[int]) -> None:
+    """T5 tool barrier: resume an IDLE request with an observation
+    (``<result>...`` token ids). Mirrors the reference
+    ``_tool_barrier``: the turn is counted; a non-empty observation
+    enters FEED (one forced token per step); an empty one goes straight
+    back to DECODE (or DONE past max_turns)."""
+    if state.mode != IDLE:
+        raise ValueError(
+            f"queue_observation in mode {state.mode} (expected IDLE)"
+        )
+    state.pending_tool = None
+    state.turn_count += 1
+    state.dirty = True
+    if token_ids:
+        state.feed_queue = list(token_ids)
+        state.mode = FEED
+    else:
+        state.text_ids_since_resume = []
+        state.in_think = False  # T4b (reference engine.py:590)
+        if state.turn_count > state.max_turns:
+            _finish(state, "max_turns")
+        else:
+            state.mode = DECODE
+
+
+def record_tool_call(state: RequestState, tag: str) -> None:
+    """T5: report an executed tool call (per-tag counter bookkeeping;
+    cap enforcement is T5's)."""
+    if tag == "python":
+        state.python_calls += 1
+    elif tag == "search":
+        state.search_calls += 1

serving/processor/test_blocked_masks.py

@@ -0,0 +1,354 @@
+#!/opt/hybridagent-vllm/venv/bin/python
+"""T4b -- tests for the reference ``blocked_token_sets`` port.
+
+Validates the token-blocking masks ported from
+togethercomputer/HybridAgent's ``HybridGenerationEngine``
+(``hybrid_agent/engine.py:140-149`` sets, ``:221/:239`` prompt
+initialization, ``:443-451`` per-step application, ``:493-496``
+think tracking, ``:434/:590`` FEED reset; ``sampling.py:75``
+masked_fill) into the T4 state machine + vLLM logits processor:
+
+  in_think  blocks {</latent>, <think>}
+  out_think blocks {</latent>, <latent>, </think>}
+
+applied to every DECODE (text) logits row, from the very first step,
+so ``</latent>`` can only come from the forced latent-exit readout
+and ``<latent>`` only inside ``<think>``. LATENT exit readouts stay
+raw; FEED/exit forcing is untouched.
+
+Run:
+  /opt/hybridagent-vllm/venv/bin/python \
+      /opt/hybridagent-vllm/processor/test_blocked_masks.py
+
+Exit code 0 = PASS, 1 = FAIL.
+"""
+from __future__ import annotations
+
+import sys
+from pathlib import Path
+
+# This file lives inside the processor/ package dir; drop that dir from
+# sys.path (the script dir) so `import processor` resolves to the
+# package itself, not processor/processor.py.
+_HERE = Path(__file__).resolve().parent
+sys.path = [p for p in sys.path if not p or Path(p).resolve() != _HERE]
+sys.path.insert(0, str(_HERE.parent))
+
+import torch
+
+import processor
+import processor.processor as pp
+from processor import HybridAgentLogitsProcessor
+from processor.state_machine import (
+    DECODE,
+    DONE,
+    FEED,
+    IDLE,
+    LATENT,
+    RequestState,
+    advance,
+    queue_observation,
+)
+from vllm.sampling_params import SamplingParams
+from vllm.v1.sample.logits_processor import BatchUpdate
+
+# Token plan (V=64): 10=<latent>, 11=</latent>, 12=<think>, 13=</think>,
+# 14=eos, 20/21=observation tokens, 40/41/30=<python>/print(1)//</python>.
+V = 64
+L, R, T, TC, EOS = 10, 11, 12, 13, 14
+PLAIN, B2 = 5, 6
+
+
+class StubTok:
+    unk_token_id = None
+    eos_token_id = EOS
+    SPECIAL = {
+        "<latent>": L,
+        "</latent>": R,
+        "<think>": T,
+        "</think>": TC,
+        "<result>": 20,
+        "</result>": 22,
+        "<python>": 40,
+        "</python>": 30,
+        "<search>": 43,
+        "</search>": 31,
+    }
+    TEXT = {5: "hello", 6: " world", 41: "print(1)"}
+
+    def __init__(self):
+        self.id2str = {**self.TEXT, **{v: k for k, v in self.SPECIAL.items()}}
+
+    def convert_tokens_to_ids(self, token):
+        return self.SPECIAL.get(token)
+
+    def decode(self, ids, skip_special_tokens=False):
+        return "".join(self.id2str.get(i, f" t{i} ") for i in ids)
+
+
+class FakeRequest:
+    def __init__(self):
+        self.pending_latent = False
+        self.engine_outputs: list[int] = []
+
+
+class FakeEngine:
+    def __init__(self, tokenizer):
+        self.tokenizer = tokenizer
+        self.requests = {"reqA": FakeRequest()}
+        self.worker_lists: dict[str, list[int]] = {"reqA": []}
+        self.inject_log: list[tuple[str, int]] = []
+
+    def get_request(self, rid):
+        return self.requests.get(rid)
+
+    def get_req_ids(self):
+        return ["reqA"]
+
+    def append(self, rid, tok):
+        self.requests[rid].engine_outputs.append(tok)
+        self.worker_lists[rid].append(tok)
+
+    def inject_token(self, rid, tok):
+        self.append(rid, tok)
+        self.inject_log.append((rid, tok))
+
+
+def make_proc(prompt_ids: list[int]):
+    pp._engine = None
+    pp._PROCESSORS.clear()
+    tok = StubTok()
+    eng = FakeEngine(tok)
+    pp.register_engine(
+        get_request=eng.get_request,
+        get_req_ids=eng.get_req_ids,
+        tokenizer=tok,
+        inject_token=eng.inject_token,
+    )
+    proc = HybridAgentLogitsProcessor(
+        vllm_config=None, device=torch.device("cpu"), is_pin_memory=False
+    )
+    params = SamplingParams(max_tokens=64, temperature=0.0)
+    out: list[int] = []
+    eng.worker_lists["reqA"] = out
+    proc.update_state(
+        BatchUpdate(
+            batch_size=1,
+            removed=(),
+            added=((0, params, list(prompt_ids), out),),
+            moved=(),
+        )
+    )
+    return proc, eng, tok
+
+
+def sm(state: RequestState, **kw):
+    defaults = dict(
+        last_token=None,
+        argmax_id=PLAIN,
+        tail_text="",
+        full_text_getter=None,
+        entry_id=L,
+        exit_id=R,
+        stop_ids=frozenset({EOS}),
+        max_tokens=None,
+        output_len=0,
+        think_open_id=T,
+        think_close_id=TC,
+    )
+    defaults.update(kw)
+    return advance(state, **defaults)
+
+
+failures: list[str] = []
+
+
+def check(name, cond, detail=""):
+    status = "ok" if cond else "FAIL"
+    print(f"  [{name}] {status}" + (f" -- {detail}" if detail and not cond else ""))
+    if not cond:
+        failures.append(f"{name}: {detail}")
+
+
+def main() -> int:
+    # ------------------------------------------------------------------
+    print("== SM1: in-think blocked set (engine.py:140-144) ==")
+    st = RequestState()
+    st.in_think = True
+    a = sm(st)
+    check("in-think blocks {</latent>, <think>}",
+          a.blocked_ids == frozenset({R, T}), str(a.blocked_ids))
+
+    print("== SM2: out-of-think blocked set (engine.py:140-144) ==")
+    st = RequestState()
+    st.in_think = False
+    a = sm(st)
+    check("out-of-think blocks {</latent>, <latent>, </think>}",
+          a.blocked_ids == frozenset({R, L, TC}), str(a.blocked_ids))
+
+    print("== SM3: mask active from the very first text step ==")
+    st = RequestState()
+    a = sm(st)  # last_token=None
+    check("first step carries a blocked set",
+          a.blocked_ids is not None, str(a.blocked_ids))
+
+    print("== SM4: in_think tracked from <think>/</think> tokens ==")
+    st = RequestState()
+    sm(st, last_token=T)
+    check("sampling <think> sets in_think", st.in_think is True, str(st.in_think))
+    sm(st, last_token=TC)
+    check("sampling </think> clears in_think", st.in_think is False,
+          str(st.in_think))
+
+    print("== SM5: FEED drain resets in_think (engine.py:434) ==")
+    st = RequestState()
+    st.mode = FEED
+    st.feed_queue = [20, 21]
+    st.in_think = True
+    a = sm(st)
+    check("feed step 1 forces token 20", a.force_id == 20, str(a.force_id))
+    check("feed step 1 has no text mask", a.blocked_ids is None,
+          str(a.blocked_ids))
+    a = sm(st)
+    check("feed drain forces last token", a.force_id == 21, str(a.force_id))
+    check("drain returns to DECODE", st.mode == DECODE, st.mode)
+    check("drain resets in_think", st.in_think is False, str(st.in_think))
+
+    print("== SM6: empty observation resets in_think (engine.py:590) ==")
+    st = RequestState()
+    st.mode = IDLE
+    st.in_think = True
+    queue_observation(st, [])
+    check("empty observation -> DECODE", st.mode == DECODE, st.mode)
+    check("in_think reset", st.in_think is False, str(st.in_think))
+
+    print("== SM7: LATENT exit readout stays raw; forcing intact ==")
+    st = RequestState()
+    st.mode = LATENT
+    st._entry_token = L
+    a = sm(st, last_token=L, argmax_id=PLAIN)
+    check("latent step has no text mask", a.blocked_ids is None,
+          str(a.blocked_ids))
+    check("latent step keeps pending_latent", a.pending_latent is True,
+          str(a.pending_latent))
+    for _ in range(3):
+        st.latent_steps += 1
+    a = sm(st, last_token=L, argmax_id=R)
+    check("anchor exit still forced at min steps",
+          a.force_id == R and a.inject_token == R, str(a))
+    check("exit clears pending_latent", a.pending_latent is False,
+          str(a.pending_latent))
+    check("exit back in DECODE", st.mode == DECODE, st.mode)
+
+    print("== SM8: IDLE/DONE carry no mask ==")
+    st = RequestState()
+    st.mode = IDLE
+    check("IDLE no mask", sm(st).blocked_ids is None)
+    st = RequestState()
+    st.mode = DONE
+    check("DONE no mask", sm(st).blocked_ids is None)
+
+    print("== SM9: unknown think tokens -> anchor-only sets ==")
+    st = RequestState()
+    st.in_think = True
+    a = sm(st, think_open_id=None, think_close_id=None)
+    check("in-think without guards = {</latent>}",
+          a.blocked_ids == frozenset({R}), str(a.blocked_ids))
+    st.in_think = False
+    a = sm(st, think_open_id=None, think_close_id=None)
+    check("out-think without guards = {</latent>, <latent>}",
+          a.blocked_ids == frozenset({R, L}), str(a.blocked_ids))
+
+    # ------------------------------------------------------------------
+    print("== P1: DECODE row masked at the blocked ids only ==")
+    proc, eng, tok = make_proc([T])  # prompt ends in <think>
+    st = processor.get_state("reqA")
+    check("in_think initialized from prompt", st.in_think is True,
+          str(st.in_think))
+    logits = torch.full((1, V), -10.0, dtype=torch.float32)
+    logits[0, R] = 0.0      # greedy argmax is the exit anchor
+    logits[0, B2] = -5.0
+    proc.apply(logits)
+    check("</latent> masked to -inf",
+          torch.isinf(logits[0, R]) and logits[0, R] < 0, str(logits[0, R]))
+    check("<think> masked to -inf (in-think guard)",
+          torch.isinf(logits[0, T]) and logits[0, T] < 0, str(logits[0, T]))
+    check("<latent> NOT blocked in-think",
+          logits[0, L] == -10.0, str(logits[0, L]))
+    check("other logits untouched",
+          logits[0, PLAIN] == -10.0 and logits[0, B2] == -5.0,
+          str(logits[0, [PLAIN, B2]]))
+    check("greedy argmax flips away from </latent>",
+          int(logits[0].argmax()) == B2, str(int(logits[0].argmax())))
+
+    print("== P2: out-of-think row blocks <latent> ==")
+    proc, eng, tok = make_proc([5])  # prompt does not end in <think>
+    st = processor.get_state("reqA")
+    check("in_think False for plain prompt", st.in_think is False,
+          str(st.in_think))
+    logits = torch.full((1, V), -10.0, dtype=torch.float32)
+    logits[0, L] = 0.0      # greedy argmax is the entry anchor
+    logits[0, B2] = -5.0
+    proc.apply(logits)
+    check("<latent> masked to -inf",
+          torch.isinf(logits[0, L]) and logits[0, L] < 0, str(logits[0, L]))
+    check("</think> masked to -inf (out-of-think guard)",
+          torch.isinf(logits[0, TC]) and logits[0, TC] < 0,
+          str(logits[0, TC]))
+    check("<think> NOT blocked out-of-think",
+          logits[0, T] == -10.0, str(logits[0, T]))
+    check("greedy argmax flips away from <latent>",
+          int(logits[0].argmax()) == B2, str(int(logits[0].argmax())))
+
+    print("== P3: FEED forcing is not weakened by the masks ==")
+    proc, eng, tok = make_proc([T])
+    # The tool tag fires on the apply() step AFTER </python> is sampled.
+    script = [40, 41, 30, PLAIN]
+    for arg in script:
+        logits = torch.full((1, V), -10.0, dtype=torch.float32)
+        logits[0, arg] = 0.0
+        proc.apply(logits)
+        eng.append("reqA", int(logits[0].argmax()))
+    st = processor.get_state("reqA")
+    check("tool tag -> IDLE with pending_tool",
+          st.mode == IDLE and st.pending_tool == ("python", "print(1)"),
+          f"{st.mode} {st.pending_tool}")
+    processor.queue_observation("reqA", [20, 21])
+    st = processor.get_state("reqA")
+    check("IDLE -> FEED", st.mode == FEED, st.mode)
+    logits = torch.full((1, V), -10.0, dtype=torch.float32)
+    proc.apply(logits)
+    n_inf = int(torch.isneginf(logits[0]).sum())
+    check("feed row forced to token 20",
+          logits[0, 20] == 0.0 and n_inf == V - 1,
+          f"tok20={logits[0, 20]} n_inf={n_inf}")
+
+    print("== P4: LATENT row logits stay raw (exit readout unmasked) ==")
+    proc, eng, tok = make_proc([T])
+    logits = torch.full((1, V), -10.0, dtype=torch.float32)
+    logits[0, L] = 0.0
+    proc.apply(logits)
+    eng.append("reqA", int(logits[0].argmax()))  # <latent> sampled
+    # Next apply: the state machine sees last_token=<latent> -> LATENT.
+    before = torch.full((1, V), -10.0, dtype=torch.float32)
+    before[0, R] = 0.0
+    logits = before.clone()
+    proc.apply(logits)
+    st = processor.get_state("reqA")
+    check("entered LATENT", st.mode == LATENT, st.mode)
+    check("latent row byte-identical (no mask, no force)",
+          torch.equal(logits, before), str(logits[0]))
+
+    # ------------------------------------------------------------------
+    print()
+    if failures:
+        print(f"FAIL ({len(failures)} problem(s)):")
+        for f in failures:
+            print(f"  - {f}")
+        return 1
+    print("PASS: T4b blocked_token_sets masks (state machine + processor)")
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

serving/serve.py

@@ -0,0 +1,176 @@
+#!/usr/bin/env python3
+"""T7b — serve the HybridAgent-8B backbone with OUR stack.
+
+This is NOT the `vllm serve` CLI. The HybridAgent latent loop needs
+IN-PROCESS engine access (the T4 state machine reaches engine-side
+`Request`s through registered handles, and the T5 orchestrator appends
+tokens via direct request mutation), so the engine runs in this process
+(`VLLM_ENABLE_V1_MULTIPROCESSING=0`) and is driven by the
+`HybridAgentOrchestrator`.
+
+T7b-full (this configuration) serves the **EXL3-quantized** backbone
+(T6: 5.0 bpw/H5, `model.safetensors` 5.98 GB) from the staged serving
+dir `serve/exl3-staged/` — symlinks to `models/HybridAgent-8B-EXL3/`
+plus the tokenizer files back-filled from the BF16 dir (stage-5 gap:
+`convert.py` did not retain them). The HybridLink `adapter.pt` loads
+as BF16 next to the checkpoint (S4). The KV budget is 15e9 B
+(~101k tokens ≈ 3 concurrent full-32k sessions). The BF16 smoke
+config (T7b-smoke) is still reachable via env:
+`HYBRIDAGENT_MODEL_DIR=/opt/hybridagent-vllm/models/HybridAgent-8B
+HYBRIDAGENT_KV_BYTES=5000000000`.
+
+Required engine semantics (feature doc L6/L8, T2 fork):
+  * VLLM_USE_V2_MODEL_RUNNER=0 — the fork modifies the V1 runner
+    (`v1/worker/gpu_model_runner.py`); vLLM 0.30 defaults to V2, where the
+    latent path would be silently inactive.
+  * async_scheduling=False — forced by the orchestrator; the async
+    scheduler keeps `num_new_tokens > 0` every step, which blocks the
+    fork's zero-token latent gate.
+  * enable_prompt_embeds=True — the fork feeds R(h) latent embeddings
+    through the mixed-batch prompt-embeds path (is_token_ids=False);
+    the T7a parity run booted with it.
+  * enforce_eager=True — the proven parity config; avoids CUDA-graph
+    memory overhead on the 24 GB 3090.
+
+KV sizing (BF16): the guide fixes the KV of the hybrid (GDN) models with
+`--kv-cache-memory-bytes` because GPU-memory profiling underestimates GDN
+layers and OOMs. HybridAgent-8B is a PLAIN full-attention transformer
+(36 × full_attention, no GDN), so utilization-based sizing would be
+acceptable; we keep the explicit `kv_cache_memory_bytes` because it makes
+the VRAM budget deterministic against the "≥100 MiB really free" battery
+criterion. BF16 weights are ~15.7 GB (vs ~5.4 GB EXL3), so the BF16 KV
+budget is 5e9 B (~33.9k tokens ≈ one full 32k session with margin);
+see README.md for the numbers.
+
+Usage (on tarninfer):
+    cd /opt/hybridagent-vllm
+    venv/bin/python serve/serve.py
+    then type one prompt per line; a BLANK line flushes (runs all pending
+    prompts in one batched loop) and prints their final texts; EOF exits.
+"""
+from __future__ import annotations
+
+import os
+import sys
+
+# ---------------------------------------------------------------------------
+# Environment — MUST be set before vllm is imported.
+# (The orchestrator module also setdefaults the two fork-semantics vars;
+#  setting them here keeps this file correct if imported standalone.)
+# ---------------------------------------------------------------------------
+os.environ.setdefault("VLLM_USE_V2_MODEL_RUNNER", "0")      # fork is V1 (L6)
+os.environ.setdefault("VLLM_ENABLE_V1_MULTIPROCESSING", "0")  # in-process core
+os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")
+os.environ.setdefault("VLLM_USE_FLASHINFER_SAMPLER", "0")
+os.environ.setdefault("CUDA_HOME", "/usr/local/cuda-13.2")
+os.environ.setdefault("VLLM_CACHE_ROOT", "/var/tmp/vllm-hybridagent")
+
+# ---------------------------------------------------------------------------
+# Configuration (T7b-full: EXL3 checkpoint; see README.md)
+# ---------------------------------------------------------------------------
+# The serving dir: EXL3 weights + adapter.pt + adapter_config.json +
+# tokenizer files, in ONE directory (the T3 backend loads the adapter from
+# next to the checkpoint). serve/exl3-staged/ is the T6 EXL3 dir
+# (models/HybridAgent-8B-EXL3) with symlinks plus the tokenizer files
+# back-filled from the BF16 dir (the stage-5 "verify tokenizer files"
+# step). Point HYBRIDAGENT_MODEL_DIR at the BF16 dir to re-run the
+# T7b-smoke config.
+MODEL_DIR = os.environ.get(
+    "HYBRIDAGENT_MODEL_DIR",
+    "/opt/hybridagent-vllm/serve/exl3-staged",
+)
+MAX_MODEL_LEN = 32_768          # the model's native context window
+# EXL3 weights are ~5.98 GB (+0.07 GB BF16 adapter), so the KV budget is
+# 15e9 B / 147,456 B-per-token ≈ 101.7k tokens ≈ 3 concurrent full-32k
+# sessions. The BF16 smoke used 5e9 (≈33.9k tokens) against its 15.7 GB
+# weights.
+KV_CACHE_MEMORY_BYTES = int(
+    os.environ.get("HYBRIDAGENT_KV_BYTES", "15_000_000_000")
+    .replace("_", "")
+)
+# Total vLLM memory budget (fraction of the 24 GB). EXL3 weights
+# (~5.98 GB) + KV (15e9 B ≈ 14.2 GB) + runtime (~2 GB) ≈ 22.2 GB →
+# 0.93 (≈22.9 GB) leaves ≥ ~100 MiB really free for the battery
+# criterion.
+GPU_MEMORY_UTILIZATION = float(os.environ.get("HYBRIDAGENT_GPU_UTIL", "0.93"))
+MAX_NUM_BATCHED_TOKENS = 1_024  # the guide's serving services use 1024
+DTYPE = "bfloat16"
+ENFORCE_EAGER = True            # proven parity config (T7a)
+ENABLE_PROMPT_EMBEDS = True     # fork latent-embed mixed-batch path
+MAX_TURNS = 30                  # the reference engine's turn cap
+MAX_TOKENS_PER_REQ = 4_096      # per-prompt output budget
+
+
+def main() -> int:
+    # The two engine-side plugins must be installed in this venv
+    # (pip install -e backend, pip install -e processor); an early import
+    # fails fast with a clear message instead of mid-boot.
+    import backend  # noqa: F401  (vllm.general_plugins entry point)
+    import processor  # noqa: F401  (vllm.logits_processors entry point)
+    # The orchestrator package is NOT installed (no entry point); the
+    # repo root must be on the path. Add it AFTER backend/processor are
+    # imported: with REPO first, `import backend` would resolve to the
+    # namespace dir /opt/hybridagent-vllm/backend/ and shadow the
+    # installed package (same ordering rule as battery/run_battery.py).
+    sys.path.insert(
+        0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
+    from orchestrator import HybridAgentOrchestrator
+
+    print(
+        f"[serve] T7b-full — model={MODEL_DIR}\n"
+        f"[serve]   max_model_len={MAX_MODEL_LEN} "
+        f"kv_cache_memory_bytes={KV_CACHE_MEMORY_BYTES} "
+        f"gpu_memory_utilization={GPU_MEMORY_UTILIZATION} "
+        f"max_num_batched_tokens={MAX_NUM_BATCHED_TOKENS} dtype={DTYPE}\n"
+        f"[serve]   VLLM_USE_V2_MODEL_RUNNER=0, async_scheduling=False (forced), "
+        f"enforce_eager={ENFORCE_EAGER}, enable_prompt_embeds={ENABLE_PROMPT_EMBEDS}"
+    )
+    if not os.path.isfile(os.path.join(MODEL_DIR, "adapter.pt")):
+        print(
+            f"[serve] WARNING: adapter.pt not found in {MODEL_DIR} — "
+            "the T3 backend will refuse to load. Fix HYBRIDAGENT_MODEL_DIR."
+        )
+
+    orch = HybridAgentOrchestrator(
+        model=MODEL_DIR,
+        max_turns=MAX_TURNS,
+        dtype=DTYPE,
+        enforce_eager=ENFORCE_EAGER,
+        enable_prompt_embeds=ENABLE_PROMPT_EMBEDS,
+        max_model_len=MAX_MODEL_LEN,
+        kv_cache_memory_bytes=KV_CACHE_MEMORY_BYTES,
+        gpu_memory_utilization=GPU_MEMORY_UTILIZATION,
+        max_num_batched_tokens=MAX_NUM_BATCHED_TOKENS,
+    )
+    print("[serve] engine up; one prompt per line, blank line = run batch")
+
+    def flush(prompts: list[str]) -> None:
+        if not prompts:
+            return
+        rids = [
+            orch.add_request(p, max_tokens=MAX_TOKENS_PER_REQ) for p in prompts
+        ]
+        orch.run()  # drives every submitted request to DONE (batched loop)
+        for rid in rids:
+            text = orch.final_text(rid)
+            print(f"\n[serve] === {rid} ===\n{text}\n", flush=True)
+
+    pending: list[str] = []
+    try:
+        for line in sys.stdin:
+            line = line.rstrip("\n")
+            if not line:
+                flush(pending)
+                pending = []
+            else:
+                pending.append(line)
+        flush(pending)
+    except KeyboardInterrupt:
+        print("\n[serve] interrupted; shutting down")
+    finally:
+        orch.shutdown()
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

serving/setup.sh

@@ -0,0 +1,51 @@
+#!/usr/bin/env bash
+# Deploy the HybridAgent-8B EXL3 model with the forked vLLM stack.
+# Run from the serving/ directory of a fresh clone of this repo:
+#   git clone https://huggingface.co/JRamirez-UAB/hybridagent-8b-exl3-vllm-experimental
+#   cd hybridagent-8b-exl3-vllm-experimental/serving
+#   ./setup.sh
+#
+# Stack: Python 3.12, vLLM 0.30.0, torch 2.13.0, exllamav3 1.5.1,
+#        OrcaSAQ2-kernel (EXL3 plugin), the hybridagent backend+processor,
+#        and the fork's v1 runner files overlaid onto the installed vLLM.
+set -euo pipefail
+cd "$(dirname "$0")"          # serving/
+
+PY=${PY:-3.12}
+VENV=${VENV:-../venv}
+ORCASAQ2_COMMIT=${ORCASAQ2_COMMIT:-7ddffb12c00c24b8bc6b1164723c04ac490a91bc}
+
+echo "[1/6] venv (python ${PY})"
+python${PY} -m venv "${VENV}"
+PIP="${VENV}/bin/pip"
+"${PIP}" install -U pip
+
+echo "[2/6] vLLM 0.30.0 (pulls torch 2.13.0)"
+"${PIP}" install "vllm==0.30.0"
+
+echo "[3/6] exllamav3 1.5.1 (cu132, torch 2.13.0, cp312)"
+W="exllamav3-1.5.1+cu132.torch2.13.0-cp312-cp312-linux_x86_64.whl"
+curl -L -o "/tmp/${W}" "https://github.com/turboderp-org/exllamav3/releases/download/v1.5.1/${W/+/%2B}"
+"${PIP}" install "/tmp/${W}"
+
+echo "[4/6] OrcaSAQ2-kernel (EXL3 plugin for vLLM)"
+rm -rf /tmp/OrcaSAQ2-kernel
+git clone https://github.com/Continuum-AI-Corp/OrcaSAQ2-kernel /tmp/OrcaSAQ2-kernel
+git -C /tmp/OrcaSAQ2-kernel checkout "${ORCASAQ2_COMMIT}"
+"${PIP}" install -e /tmp/OrcaSAQ2-kernel
+
+echo "[5/6] hybridagent engine-side plugins (backend + processor)"
+"${PIP}" install -e ./backend -e ./processor
+
+echo "[6/6] overlay the fork onto the installed vLLM v1 runner"
+SP="$("${VENV}/bin/python" -c 'import vllm, os; print(os.path.dirname(vllm.__file__))')"
+cp -v fork/v1/request.py                    "${SP}/v1/request.py"
+cp -v fork/v1/worker/gpu_model_runner.py    "${SP}/v1/worker/gpu_model_runner.py"
+cp -v fork/v1/core/kv_cache_manager.py      "${SP}/v1/core/kv_cache_manager.py"
+cp -v fork/v1/core/sched/output.py          "${SP}/v1/core/sched/output.py"
+cp -v fork/v1/core/sched/scheduler.py       "${SP}/v1/core/sched/scheduler.py"
+
+echo
+echo "Setup complete. Serve with (MODEL_DIR = the repo root, which holds the weights):"
+echo "  cd .."
+echo "  HYBRIDAGENT_MODEL_DIR=\$(pwd) ${VENV}/bin/python serving/serve.py"