vllm.entrypoints.generate.structured_decisions.strategies

Read strategies: how a model is asked for label probabilities.

Only models in NEXT_TOKEN_ARCHITECTURES are currently supported.

Classes:

• NextTokenStrategy –

  Autoregressive models. Each question is one request: the state, then the

• ReadStrategy –

Functions:

• reply_label_ids –

  The prompt's ids after its last added token, and the token each of LABELS

• select_read_strategy –

  Raises ValueError when the model cannot serve structured decisions. The

NextTokenStrategy

Bases: ReadStrategy

Autoregressive models. Each question is one request: the state, then the question with its labeled options, and the logprobs of the label tokens as the reply's first token. The requests share the state. With prefix caching, it is prefilled once.

Source code in vllm/entrypoints/generate/structured_decisions/strategies.py

class NextTokenStrategy(ReadStrategy):
    """Autoregressive models. Each question is one request: the state, then the
    question with its labeled options, and the logprobs of the label tokens as
    the reply's first token. The requests share the state. With prefix caching,
    it is prefilled once."""

    def __init__(self, context: ReadContext):
        super().__init__(context)
        tokenizer = context.online_renderer.renderer.get_tokenizer()
        probe = tokenizer.apply_chat_template(
            [{"role": "user", "content": "x"}],
            **{
                **context.default_chat_template_kwargs,
                "chat_template": context.chat_template,
                "add_generation_prompt": True,
                "tokenize": True,
                "return_dict": False,
                "enable_thinking": False,
            },
        )
        if isinstance(probe, str):
            probe = tokenizer.encode(probe, add_special_tokens=False)
        # Every prompt ends with the same generation prompt, so the labels'
        # tokens are the same for every question.
        self.tail, self.label_ids = reply_label_ids(tokenizer, probe)

    def limits(self) -> DecisionLimits:
        return DecisionLimits(max_questions=64, max_options=len(LABELS))

    def _read_request(
        self, chat_template_kwargs: dict[str, Any] | None
    ) -> ReadPromptRequest:
        if (chat_template_kwargs or {}).get("enable_thinking"):
            raise StructuredDecisionError(
                "a read is the reply's first token, so thinking must be off"
            )
        return ReadPromptRequest(chat_template_kwargs=chat_template_kwargs)

    async def _render(
        self, read_request: ReadPromptRequest, messages: list[dict[str, Any]]
    ) -> tuple[EngineInput, list[int]]:
        ctx = self.context
        _, (engine_input,) = await ctx.online_renderer.preprocess_chat(
            read_request,
            messages,
            default_template=ctx.chat_template,
            default_template_content_format=ctx.chat_template_content_format,
            default_template_kwargs=ctx.default_chat_template_kwargs,
        )
        prompt_ids = extract_prompt_components(
            ctx.engine_client.model_config, engine_input
        ).token_ids
        return engine_input, list(prompt_ids or [])

    async def read(
        self,
        questions: list[Question],
        instructions: str | None,
        state: str,
        *,
        request_id: str,
        chat_template_kwargs: dict[str, Any] | None,
        lora_request: LoRARequest | None,
        priority: int,
    ) -> list[QuestionRead]:
        ctx = self.context
        read_request = self._read_request(chat_template_kwargs)

        slots, engine_inputs = [], []
        for q in questions:
            messages = [{"role": "user", "content": f"{state}\n\n{q.type.prompt(q)}"}]
            if instructions:
                messages.insert(0, {"role": "system", "content": instructions})
            engine_input, prompt_ids = await self._render(read_request, messages)
            if prompt_ids[len(prompt_ids) - len(self.tail) :] != self.tail:
                raise StructuredDecisionError(
                    "these chat options end the prompt differently, so the "
                    "labels' tokens are unknown"
                )
            slots.append(self.label_ids[: len(q.labels)])
            engine_inputs.append(engine_input)

        label_reads = await next_token_label_reads(
            ctx.engine_client,
            engine_inputs,
            [
                SamplingParams(max_tokens=1, temperature=0.0, logprob_token_ids=ids)
                for ids in slots
            ],
            request_id,
            lora_request=lora_request,
            priority=priority,
        )

        reads = []
        for ids, label_read in zip(slots, label_reads):
            output = label_read.result.outputs[0]
            reads.append(
                QuestionRead(
                    probs=label_softmax(label_read.logprobs),
                    label_mass=sum(math.exp(lp) for lp in label_read.logprobs),
                    argmax_is_label=bool(output.token_ids)
                    and output.token_ids[0] in ids,
                    input_tokens=len(label_read.result.prompt_token_ids or ()),
                    output_tokens=len(output.token_ids),
                )
            )
        return reads

ReadStrategy

Bases: ABC

Methods:

• read –

  One read per question, in order.

Source code in vllm/entrypoints/generate/structured_decisions/strategies.py

class ReadStrategy(ABC):
    def __init__(self, context: ReadContext):
        self.context = context

    @abstractmethod
    def limits(self) -> DecisionLimits: ...

    @abstractmethod
    async def read(
        self,
        questions: list[Question],
        instructions: str | None,
        state: str,
        *,
        request_id: str,
        chat_template_kwargs: dict[str, Any] | None,
        lora_request: LoRARequest | None,
        priority: int,
    ) -> list[QuestionRead]:
        """One read per question, in order."""

read(questions, instructions, state, *, request_id, chat_template_kwargs, lora_request, priority) abstractmethod async

One read per question, in order.

Source code in vllm/entrypoints/generate/structured_decisions/strategies.py

@abstractmethod
async def read(
    self,
    questions: list[Question],
    instructions: str | None,
    state: str,
    *,
    request_id: str,
    chat_template_kwargs: dict[str, Any] | None,
    lora_request: LoRARequest | None,
    priority: int,
) -> list[QuestionRead]:
    """One read per question, in order."""

reply_label_ids(tokenizer, prompt_ids)

The prompt's ids after its last added token, and the token each of LABELS adds after them as the first token of the reply. Raises ValueError if a label is not one distinct token there.

Source code in vllm/entrypoints/generate/structured_decisions/strategies.py

def reply_label_ids(
    tokenizer: TokenizerLike, prompt_ids: Sequence[int]
) -> tuple[list[int], list[int]]:
    """The prompt's ids after its last added token, and the token each of LABELS
    adds after them as the first token of the reply. Raises ValueError if a
    label is not one distinct token there."""
    added = set(tokenizer.get_added_vocab().values())
    start = max((i + 1 for i, t in enumerate(prompt_ids) if t in added), default=0)
    tail = list(prompt_ids[start:])
    text = tokenizer.decode(tail)
    ids: list[int] = []
    for label in LABELS:
        extended = tokenizer.encode(text + label, add_special_tokens=False)
        if extended[:-1] != tail or extended[-1] in ids:
            raise ValueError(
                f"label {label!r} is not one distinct token after this model's "
                "chat prompt"
            )
        ids.append(extended[-1])
    return tail, ids

select_read_strategy(model_config)

Raises ValueError when the model cannot serve structured decisions. The server opted in with --enable-structured-decisions, so startup fails.

Source code in vllm/entrypoints/generate/structured_decisions/strategies.py

def select_read_strategy(model_config: ModelConfig) -> type[ReadStrategy]:
    """Raises ValueError when the model cannot serve structured decisions. The
    server opted in with --enable-structured-decisions, so startup fails."""
    if model_config.architecture not in NEXT_TOKEN_ARCHITECTURES:
        raise ValueError(
            "--enable-structured-decisions does not support "
            f"{model_config.architecture}. Supported architectures: "
            f"{sorted(NEXT_TOKEN_ARCHITECTURES)}"
        )
    if model_config.logprobs_mode not in LOGPROBS_MODES:
        raise ValueError(
            "--enable-structured-decisions needs --logprobs-mode raw_logprobs or "
            f"processed_logprobs, not {model_config.logprobs_mode}"
        )
    return NextTokenStrategy