Fix logprobs issues.

[dcampora]

Fix logprobs issues

[dcampora]

/bot run

[tensorrt-cicd]

PR_Github #8548 [ run ] triggered by Bot

[vegaluisjose]

yes, this seems to work now.

import torch

from tensorrt_llm._torch import LLM
from tensorrt_llm.llmapi import KvCacheConfig
from tensorrt_llm.sampling_params import SamplingParams


def main():
    model_dir = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"

    text_prompts = [
        "The future of AI is",
        "The president of the United States is",
    ]

    num_prompts = len(text_prompts)

    llm = LLM(
        model=model_dir,
        max_batch_size=num_prompts,
        use_cuda_graph=False,
        kv_cache_config=KvCacheConfig(enable_block_reuse=False),
        enable_trtllm_sampler=True,
    )

    sampling_params = SamplingParams(
        max_tokens=9,
        temperature=0.0,
        add_special_tokens=False,
        logprobs=True,
    )

    outputs = llm.generate(text_prompts, sampling_params)

    for out in outputs:
        print("Logprobs:", out.outputs[0].logprobs)
        print("Text:", out.outputs[0].text)


if __name__ == "__main__":
    main()

Produces

Logprobs: [{11785: Logprob(logprob=-1.8318084478378296, rank=1)}, {29892: Logprob(logprob=-0.6221925020217896, rank=1)}, {322: Logprob(logprob=-0.4919593930244446, rank=1)}, {372: Logprob(logprob=-1.5133262872695923, rank=1)}, {29915: Logprob(logprob=-0.6661361455917358, rank=1)}, {29879: Logprob(logprob=-0.012634724378585815, rank=1)}, {5566: Logprob(logprob=-1.5086636543273926, rank=1)}, {11407: Logprob(logprob=-4.2439409298822284e-05, rank=1)}, {304: Logprob(logprob=-0.014878500252962112, rank=1)}]
Text: bright, and it's exciting to
Logprobs: [{263: Logprob(logprob=-1.745693325996399, rank=1)}, {5829: Logprob(logprob=-3.1098272800445557, rank=1)}, {310: Logprob(logprob=-0.2692946195602417, rank=1)}, {278: Logprob(logprob=-1.0950253009796143, rank=1)}, {3082: Logprob(logprob=-1.4572322368621826, rank=1)}, {12561: Logprob(logprob=-1.1409716606140137, rank=1)}, {29889: Logprob(logprob=-0.8791855573654175, rank=1)}, {940: Logprob(logprob=-1.2281692028045654, rank=1)}, {338: Logprob(logprob=-1.3323622941970825, rank=1)}]
Text: a symbol of the American dream. He is

[vegaluisjose]

fyi, logprobs still need this fix. @tomeras91 @amitz-nv

[Funatiq]

Why was this not caught by the test? Is there anything we could change/add to improve the test?

[tensorrt-cicd]

PR_Github #8548 [ run ] completed with state SUCCESS
/LLM/main/L0_MergeRequest_PR pipeline #6199 completed with status: 'SUCCESS'