Is this page helpful?

Custom Logits Processing#

NIM LLM and VLM supports custom logits processing by passing through to vLLM’s native logits processor API. You can write your own processor, mount it into the container as a volume, and enable it with a single CLI flag.

Custom processors operate at the batch level. vLLM batches multiple requests into a single tensor for GPU efficiency, so your processor receives the entire batch together. Each row in the tensor is one request.

You can still toggle behavior per request. Clients pass custom arguments in vllm_xargs. Your processor reads each row’s extra_args and modifies only the rows that opt in.

Before you start, complete the following prerequisites:

  1. Write a processor that extends vllm.v1.sample.logits_processor.LogitsProcessor, or use the example in this procedure.

  2. Confirm that the processor module is importable from the PYTHONPATH you will set on the container.

To run a custom logits processor, complete the following steps:

  1. Implement the five LogitsProcessor methods in your processor script.

    Custom processors extend vllm.v1.sample.logits_processor.LogitsProcessor. The following table describes each method:

    Method

    Purpose

    validate_params

    Class method. Validates per-request parameters when a request arrives. Raise ValueError to reject the request.

    __init__

    Runs during server startup. Initialize any state your processor needs.

    apply

    Called every engine step. Receives the full logits tensor (num_requests, vocab_size) and returns the modified tensor.

    is_argmax_invariant

    Return True if your processor never changes which token has the highest logit. Return False otherwise.

    update_state

    Called every engine step with a BatchUpdate describing which requests were added, removed, or moved in the batch. Use this to track per-request state.
    Example: BanTokenProcessor

    The following processor bans a single token per request. Requests opt in by sending {"vllm_xargs": {"ban_token_id": <int>}}. Requests that do not send this parameter are unaffected.

    import torch
from vllm.config import VllmConfig
from vllm.sampling_params import SamplingParams
from vllm.v1.sample.logits_processor import BatchUpdate, LogitsProcessor, MoveDirectionality


class BanTokenProcessor(LogitsProcessor):

    @classmethod
    def validate_params(cls, params: SamplingParams):
        ban = params.extra_args and params.extra_args.get("ban_token_id")
        if ban is not None and not isinstance(ban, int):
            raise ValueError(f"ban_token_id must be int, got {type(ban).__name__}")

    def __init__(self, vllm_config: VllmConfig, device: torch.device, is_pin_memory: bool):
        self.banned: dict[int, int] = {}

    def is_argmax_invariant(self) -> bool:
        return False

    def update_state(self, batch_update: BatchUpdate | None) -> None:
        if not batch_update:
            return

        for index, params, _, _ in batch_update.added:
            ban = params.extra_args and params.extra_args.get("ban_token_id")
            if ban is not None:
                self.banned[index] = ban
            else:
                self.banned.pop(index, None)

        if not self.banned:
            return

        for index in batch_update.removed:
            self.banned.pop(index, None)

        for src, dst, direction in batch_update.moved:
            src_val = self.banned.pop(src, None)
            dst_val = self.banned.pop(dst, None)
            if src_val is not None:
                self.banned[dst] = src_val
            if direction == MoveDirectionality.SWAP and dst_val is not None:
                self.banned[src] = dst_val

    def apply(self, logits: torch.Tensor) -> torch.Tensor:
        if not self.banned:
            return logits

        for req_idx, token_id in self.banned.items():
            logits[req_idx, token_id] = float("-inf")

        return logits

  2. Place your processor script in a local directory, for example /home/user/my_processors/ban_token.py.

  3. Volume-mount that directory into the container.

  4. Pass the processor’s module path to NIM by using the --logits-processors CLI argument.

    docker run --gpus all \
-v /home/user/my_processors:/opt/nim/my_processors \
-e PYTHONPATH=/opt/nim \
-e NIM_MODEL_PATH=hf://meta-llama/Llama-3.2-1B-Instruct \
-p 8000:8000 \
nim-llm:local \
nim-serve --logits-processors my_processors.ban_token:BanTokenProcessor

    The --logits-processors value uses Python’s module.submodule:ClassName format. The code must be importable from PYTHONPATH.

  5. Optional: Enable the processor for a specific request by passing custom arguments in vllm_xargs.

    Custom logits processors operate at the batch level, but clients can enable specific logits processors per request by passing custom arguments using vllm_xargs.

    curl -X POST http://localhost:8000/v1/completions \
  -H "Content-Type: application/json" \
  -d '{
    "model": "meta-llama/Llama-3.2-1B-Instruct",
    "prompt": "Once upon a time...",
    "max_tokens": 64,
    "vllm_xargs": {"ban_token_id": 42}
  }'

    For more details on invoking custom logits processors per request, refer to the vLLM documentation.

Next Steps#

After the processor loads, continue with Custom Parsers and Chat Templates if you also need a custom tool-call or reasoning parser.

On this page