How to obtain the logits of LLM

General
21

Yes, I always need the code example! Thanks!

22

Certainly! Here’s how you can introduce and use an is_sigmoid parameter in vLLM V1 using extra_args, following the current extensibility pattern (Custom sampling params RFC):

1. When creating SamplingParams, add is_sigmoid to extra_args:

from vllm import SamplingParams

sampling_params = SamplingParams(
    temperature=0.8,
    top_p=0.95,
    extra_args={"is_sigmoid": True}
)

2. In your custom sampler logic (e.g., in vllm/v1/sample/sampler.py), access is_sigmoid from extra_args:

import torch

class Sampler(nn.Module):
    ...
    def forward(
        self,
        logits: torch.Tensor,
        sampling_metadata,
        sampling_params,  # Make sure to pass this in your call chain
        # ... other args
    ):
        is_sigmoid = sampling_params.extra_args.get("is_sigmoid", False)
        if is_sigmoid:
            probs = torch.sigmoid(logits)
            # Use or return probs as needed
            return {"sigmoid_probs": probs}
        else:
            logprobs = torch.log_softmax(logits, dim=-1)
            return {"logprobs": logprobs}

Note:

  • You must ensure that the relevant SamplingParams instance is available in your sampler logic. If not, you may need to modify the call chain to pass it.
  • This approach is forward-compatible with vLLM’s planned plugin and custom argument support.

Would you like more detail on how to wire SamplingParams into the V1 sampler if it’s not already available?

Sources: