Probability / logit of output tokens

[Simratpreet]

I am trying to log the probability/logit of output token in top K sampling kernel. This is the place where I am trying to log.

My understanding is that outputLogProbs should contain the probability but it is always 0. Also my understanding is sSum variable should contain sum of all logits, but it is always equal to expLogit.

here is actual output for one of the selected tokens:
outputLogProbs[curSeqLen * maxBatchSize + batchSlot]= 0.000000
sSum: 0.000201, logSum: -8.511453, logprob: -8.511453, explogit: 0.000201

In summary:

1. Is this the right place to get the final logit of the token in Top K sampling mode
2. Why would the outputLogprobs be zero and sSum match expLogit? Is this expected?

[Advait251206]

Yes, this is expected behavior given how the Top-K sampling kernel is implemented.

The key point is that the kernel performs Top-K filtering before sampling. By the time execution reaches the sampling stage, only the Top-K candidates remain, and the normalization is performed over those candidates—not over the entire vocabulary.

Regarding your observations:

outputLogProbs is always 0.

This is normal unless outputLogProbs has been explicitly enabled and written by the runtime. In many TensorRT-LLM configurations, this buffer is optional and may remain zero-initialized if log-prob collection is disabled or deferred to another stage. Simply seeing a non-null pointer doesn't necessarily mean it has been populated.

sSum always equals expLogit.

This typically means the sampled token is the only surviving candidate after filtering. In that case,

sSum = exp(logit_selected)

and therefore

log_prob = log(exp(logit_selected) / sSum)
         = log(1)
         = 0

or equivalently,

logSum == log(expLogit)

which matches your observation:

sSum     = 0.000201
expLogit = 0.000201
logSum   = -8.511453
logProb  = -8.511453

Since

log(expLogit) = -8.511453

everything is internally consistent.

If k > 1, however, then you would normally expect:

sSum > expLogit

because the denominator should include the exponentials of all retained Top-K candidates.

So a few questions that may help narrow this down:

• What value of Top-K are you using? Is it actually greater than 1?
• Are you also using Top-P, Min-P, or any other filtering that could reduce the candidate set to a single token?
• Are you trying to obtain the raw model logits before sampling, or the final normalized log probability after Top-K filtering? Those are different quantities, and they are computed at different stages of the decoding pipeline.

If your goal is to obtain the model's original logits over the full vocabulary, the Top-K sampling kernel is probably not the ideal place, since the logits have already been filtered by then. It would be better to inspect them immediately after the model forward pass and before any sampling or filtering kernels are applied.

---

If this answer helped or pointed you in the right direction, I'd appreciate it if you could mark it as the accepted answer so it's easier for others with the same issue to find.

Also, if you found my contribution useful, I'd appreciate it if you could check out my GitHub profile, follow me, and star any repositories you find interesting.

GitHub: https://github.com/Advait251206