[WIP] Support logprob calculation for loglikelihood approach

[vvchernov]

Loglikelihood approach is needed to check accuracy of a model on popular datasets and tasks like MMLU and BigBench. Particularly HuggingFace leaderboard bases on such tasks only.

Notes:

1. The branch bases on the work branch of Enable Logprobs in MLC Batch Serving #82, the latter PR should be merged before it.
2. Loglikelihood requires sequence of logits after prefill-like inference. But in mlc-llm on the model topology side the last logits set is splitted from all. I've added all logits to output tuple on llama topology side for multi-batch implementation.

[masahi]

Why do we need this for batch-serving? For batched inference, we use PyTorch for sampling. So adding logprob support is easy.
https://github.com/octoml/mlc-llm/blob/batch-serving/serve/mlc_serve/model/paged_cache_model.py#L246-L319

[vvchernov]

Hello @masahi! Thank you for your reference. The main idea is support accuracy benchmark of octoml endpoints on tasks (like MMLU, HellaSwag) with loglikelihood approach. Unfortunately I'm not familiar with serve implementation and made mistake when implemented it on "old" part of mlc-llm. I plan to use this PR and do it on serve side. Possibly logprobs calculation has been already done or can be easily done, but it also needs some high-level API for request-response of logprobs.

[sunggg]

Since we have #82, do we still need this PR?

[vvchernov]

Hello @sunggg! Yes, of course. Functionality from #82 allows to get logprobs info for new generated tokens. This PR allows to get logprobs for all tokens from input prompt (prefill step) which used for loglikelihood calculation. Particularly here I modify Relax model due to it cut the last set of logits, but I need all of them.
I'm helping with #82 due to want it will be merged first. My branch bases on the branch from the PR and contains the code from there, but I rebase it when the PR is merged.

[vvchernov]

Close due to transfer to octoml/mlc-serve/pull/56