vllm.model_executor.models.unlimited_ocr2

Optional page-count stopping guard for Unlimited-OCR-2.

Classes:

• PageCountGuard –

  Per-request page-count state machine over the generated tokens.

• PageCountGuardLogitsProcessor –

  Page-count guard for Unlimited-OCR-2 multi-page parsing.

PageCountGuard

Per-request page-count state machine over the generated tokens.

CPU-friendly core of the processor: token bookkeeping is plain Python ints/lists, only the one-row logits edit touches torch.

Methods:

• seed_expected_pages –

  Count the prompt's page markers; called once the prompt is known.

• step –

  Apply the guard to one row of logits, in place.

Source code in vllm/model_executor/models/unlimited_ocr2.py

class PageCountGuard:
    """Per-request page-count state machine over the generated tokens.

    CPU-friendly core of the processor: token bookkeeping is plain Python
    ints/lists, only the one-row logits edit touches torch.
    """

    def __init__(
        self,
        page_token_id: int,
        eos_token_id: int,
        max_page_tokens: int = DEFAULT_MAX_PAGE_TOKENS,
    ):
        self.page_token_id = page_token_id
        self.eos_token_id = eos_token_id
        self.max_page_tokens = max_page_tokens
        self.expected_pages = -1
        self.marker_cum = 0
        self.page_len = 0
        self.scanned = 0

    def seed_expected_pages(self, prompt_ids: list[int]) -> None:
        """Count the prompt's page markers; called once the prompt is known."""
        self.expected_pages = prompt_ids.count(self.page_token_id)

    def step(self, output_ids: list[int], logits: torch.Tensor) -> None:
        """Apply the guard to one row of logits, in place."""
        if self.expected_pages <= 0:
            return
        n = len(output_ids)
        if n > self.scanned:
            tail = output_ids[self.scanned :]
            num_markers = tail.count(self.page_token_id)
            if num_markers:
                self.marker_cum += num_markers
                self.page_len = tail[::-1].index(self.page_token_id)
            else:
                self.page_len += len(tail)
            self.scanned = n

        if self.marker_cum > self.expected_pages:
            _force_token(logits, self.eos_token_id)
        elif 0 < self.max_page_tokens <= self.page_len:
            force = (
                self.page_token_id
                if self.marker_cum < self.expected_pages
                else self.eos_token_id
            )
            _force_token(logits, force)
        elif self.marker_cum < self.expected_pages:
            finite = torch.isfinite(logits)
            finite[self.eos_token_id] = False
            if finite.any():
                logits[self.eos_token_id] = -float("inf")

seed_expected_pages(prompt_ids)

Count the prompt's page markers; called once the prompt is known.

Source code in vllm/model_executor/models/unlimited_ocr2.py

def seed_expected_pages(self, prompt_ids: list[int]) -> None:
    """Count the prompt's page markers; called once the prompt is known."""
    self.expected_pages = prompt_ids.count(self.page_token_id)

step(output_ids, logits)

Apply the guard to one row of logits, in place.

Source code in vllm/model_executor/models/unlimited_ocr2.py

def step(self, output_ids: list[int], logits: torch.Tensor) -> None:
    """Apply the guard to one row of logits, in place."""
    if self.expected_pages <= 0:
        return
    n = len(output_ids)
    if n > self.scanned:
        tail = output_ids[self.scanned :]
        num_markers = tail.count(self.page_token_id)
        if num_markers:
            self.marker_cum += num_markers
            self.page_len = tail[::-1].index(self.page_token_id)
        else:
            self.page_len += len(tail)
        self.scanned = n

    if self.marker_cum > self.expected_pages:
        _force_token(logits, self.eos_token_id)
    elif 0 < self.max_page_tokens <= self.page_len:
        force = (
            self.page_token_id
            if self.marker_cum < self.expected_pages
            else self.eos_token_id
        )
        _force_token(logits, force)
    elif self.marker_cum < self.expected_pages:
        finite = torch.isfinite(logits)
        finite[self.eos_token_id] = False
        if finite.any():
            logits[self.eos_token_id] = -float("inf")

PageCountGuardLogitsProcessor

Bases: LogitsProcessor

Page-count guard for Unlimited-OCR-2 multi-page parsing.

In the multi-page layout ("<PAGE><image>" * n_pages + "Multi page parsing.") the model must emit exactly one <PAGE> marker per prompt page marker. This request-level guard enforces that contract, mirroring the reference _uocr_page_count_guard:

• while fewer <PAGE> markers than prompt pages have been generated, EOS is masked so a page cannot be terminated early;
• once the page count is reached EOS is allowed, and generating further <PAGE> markers forces EOS;
• each page is capped at max_page_tokens generated tokens -- hitting the cap forces <PAGE> (next page) or EOS (last page);
• prompts without a <PAGE> marker (single-page requests) pass through.

Forced transitions yield to existing token masks from other sampling constraints when their target token is unavailable. EOS is retained if suppressing it would leave no finite logits.

Register it via the CLI::

--logits-processors \
  vllm.model_executor.models.unlimited_ocr2:PageCountGuardLogitsProcessor

Per-request parameters are read from SamplingParams.extra_args (page_token_id, eos_token_id, max_page_tokens); the token ids fall back to hf_config.sparse_attention_config.page_token_id and hf_config.eos_token_id when not given. The processor is inactive for requests (or v1 configs) where no page token id can be resolved.

GPU state is keyed by the persistent request slot. The prompt is scanned once; subsequent steps scan only newly committed output tokens.

Source code in vllm/model_executor/models/unlimited_ocr2.py

class PageCountGuardLogitsProcessor(LogitsProcessor):
    """Page-count guard for Unlimited-OCR-2 multi-page parsing.

    In the multi-page layout (``"<PAGE><image>" * n_pages + "Multi page
    parsing."``) the model must emit exactly one ``<PAGE>`` marker per prompt
    page marker. This request-level guard enforces that contract, mirroring
    the reference ``_uocr_page_count_guard``:

    * while fewer ``<PAGE>`` markers than prompt pages have been generated,
      EOS is masked so a page cannot be terminated early;
    * once the page count is reached EOS is allowed, and generating further
      ``<PAGE>`` markers forces EOS;
    * each page is capped at ``max_page_tokens`` generated tokens -- hitting
      the cap forces ``<PAGE>`` (next page) or EOS (last page);
    * prompts without a ``<PAGE>`` marker (single-page requests) pass
      through.

    Forced transitions yield to existing token masks from other sampling
    constraints when their target token is unavailable. EOS is retained if
    suppressing it would leave no finite logits.

    Register it via the CLI::

        --logits-processors \\
          vllm.model_executor.models.unlimited_ocr2:PageCountGuardLogitsProcessor

    Per-request parameters are read from ``SamplingParams.extra_args``
    (``page_token_id``, ``eos_token_id``, ``max_page_tokens``); the token ids
    fall back to ``hf_config.sparse_attention_config.page_token_id`` and
    ``hf_config.eos_token_id`` when not given. The processor is inactive for
    requests (or v1 configs) where no page token id can be resolved.

    GPU state is keyed by the persistent request slot. The prompt is scanned
    once; subsequent steps scan only newly committed output tokens.
    """

    def __init__(self, vllm_config: VllmConfig, req_states: LogitsProcRequestState):
        from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor

        if vllm_config.speculative_config is not None:
            raise ValueError("Page count guard does not support speculative decoding")
        hf_config = vllm_config.model_config.hf_config
        sparse_config = getattr(hf_config, "sparse_attention_config", None)
        self._default_page_token_id = self._read(sparse_config, "page_token_id")
        eos_token_id = getattr(hf_config, "eos_token_id", None)
        if isinstance(eos_token_id, list):
            eos_token_id = eos_token_id[0] if eos_token_id else None
        self._default_eos_token_id = eos_token_id

        self.all_token_ids = req_states.all_token_ids
        self.prompt_len_np = req_states.prompt_len.np
        self._prompt_len = req_states.prompt_len
        self._params = StagedWriteTensor(
            (req_states.max_num_reqs, 3), torch.int64, req_states.device
        )
        self._state = StagedWriteTensor(
            (req_states.max_num_reqs, 4), torch.int64, req_states.device
        )
        self._guards: dict[int, PageCountGuard] = {}

    @staticmethod
    def _read(config: Any, key: str) -> Any:
        if config is None:
            return None
        if isinstance(config, Mapping):
            return config.get(key)
        return getattr(config, key, None)

    @staticmethod
    def _validate_token_id(name: str, value: Any):
        if (
            not isinstance(value, int)
            or isinstance(value, bool)
            or not 0 <= value < 2**63
        ):
            raise ValueError(
                f"`{name}` has to be a non-negative int64 integer, got {value}."
            )

    @classmethod
    def validate_params(cls, sampling_params: SamplingParams) -> None:
        extra_args = sampling_params.extra_args or {}
        for name in ("page_token_id", "eos_token_id"):
            if name in extra_args:
                value = extra_args[name]
                if name == "page_token_id" and value is None:
                    continue
                cls._validate_token_id(name, value)
        if "max_page_tokens" in extra_args:
            value = extra_args["max_page_tokens"]
            if (
                not isinstance(value, int)
                or isinstance(value, bool)
                or not -(2**63) <= value < 2**63
            ):
                raise ValueError(
                    f"`max_page_tokens` has to be an int64 integer, got {value}."
                )

    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> bool:
        self.validate_params(sampling_params)
        self._guards.pop(req_idx, None)
        self._state.stage_write(req_idx, 0, [-1, 0, 0, -1])
        extra_args = sampling_params.extra_args or {}
        page_token_id = extra_args.get("page_token_id", self._default_page_token_id)
        if page_token_id is None:
            self._params.stage_write(req_idx, 0, [-1, 0, 0])
            return False
        self._validate_token_id("page_token_id", page_token_id)

        eos_token_id = extra_args.get("eos_token_id", self._default_eos_token_id)
        if eos_token_id is None:
            raise ValueError(
                "`eos_token_id` is required for the page count guard but could "
                "not be resolved from `extra_args` or the model config."
            )
        self._validate_token_id("eos_token_id", eos_token_id)

        max_page_tokens = extra_args.get("max_page_tokens", DEFAULT_MAX_PAGE_TOKENS)
        if not isinstance(max_page_tokens, int) or isinstance(max_page_tokens, bool):
            raise ValueError(
                f"`max_page_tokens` has to be an integer, got {max_page_tokens}."
            )

        self._guards[req_idx] = PageCountGuard(
            page_token_id, eos_token_id, max_page_tokens
        )
        self._params.stage_write(
            req_idx, 0, [page_token_id, eos_token_id, max_page_tokens]
        )
        return True

    def apply_staged_writes(self) -> None:
        self._params.apply_write()
        self._state.apply_write()

    def apply(self, logits: torch.Tensor, ctx: LogitsContext) -> torch.Tensor:
        if not self._guards:
            return logits
        _page_count_guard_kernel[(logits.shape[0],)](
            logits,
            self.all_token_ids.gpu,
            self._prompt_len.gpu,
            ctx.pos,
            ctx.expanded_idx_mapping,
            self._params.gpu,
            self._state.gpu,
            logits.stride(0),
            logits.stride(1),
            self.all_token_ids.gpu.stride(0),
            logits.shape[1],
            1024,
        )
        return logits