[Question] Understanding Generation Logits & Context Logits

[here4dadata]

After an engine is built with the --gather_all_token_logits and a call is made through the backend with return_context_logits: True, return_generation_logits: True, It seems as though to piece together the full text_output, we have to grab the very last context_logits, and the remaining generation_logits.

An example would be prompting a model like Llama-3 to "Say Yes!" ... The reply from the model being : "YES!"

Imagine calling the /generate endpoint of ensemble with a request that contains this in the data body:

data = {
    "text_input": f'{prompt}',
    "max_tokens": 8,
    "return_log_probs": True,
    "return_context_logits": True,
    "return_generation_logits": True, 
    "parameters": {
       "temperature": 1,
       "top_k": 5,
    }
}

text_output is : "YES!"

Now when you take the np.argmax of the generation_logits and feed it back through the tokenizer, all you would see is:

• index 0 : (token_id: 0, token: "!")
• index 1 : (token_id: 128009, token: '<|eot_id|>')

The token associated with "YES" is actually found when you np.argmax the last result in the context_logits matrix, i.e:

• index -1 : (token_id: 14331, token: "YES")

Should the "YES" not be in the generation_logits as well? Am I not understanding something fundamental?

I am using TensorRT-LLM v0.10.0 and the nvcr.io/nvidia/tritonserver:24.06-trtllm-python-py3 image to host the model.

Thanks for any help!

[byshiue]

That's because the "YES" token is computed by logits of context phase. So, it is stored in the context logits.

[here4dadata]

Thanks for that clarification @byshiue.

[VsonicV]

@byshiue I think the current implementation is problematic. In the example here, the "YES" is the first token of the generated response (This is not the "Yes" in the prompt.), so the logits for this token should be included in generation_logits, not in context_logits. The current implementation is very inconvenient if the user simply wants to get all the logits of the generated sequence, i.e., they have to get the last element of the context_logits as the logits of the first token of the generation. According to the official documentation, context_logit should simply be "Context logits for input", while generation_logits be "Generation logits for each output".

[here4dadata]

Re-opening based on @VsonicV 's post