Update logprob for prefill step

[vvchernov]

See PR for details

[zxybazh]

Hi thanks for sending in the PR! The mlc serve part fix looks pretty good to me. On the relax model part, would you please explain a bit more about all_logits and the use of last_logits to help me understand the usage in logprobs?

[vvchernov]

I think some misunderstanding appeared due to it has not been done yet.
The main idea is that we should output all logits for logprobs from prefill step, but initially it is always cut on topology side.
Two practice scenarios I see for utilization of it: 1. loglikelihood approach of accuracy benchmark 2. speculative decoding

Definitions:
last_logits is logits of new token (last token in the sequence). It is that output now without fix
logits is all logits of token sequence, they are output only with all_logits key

I see two options how it can be implemented, both have its pons and cons:

1. Return logits instead of last_logits. It is back-compatible, but we need to follow and cut last logits (token) on cpu side outside of topology which potentially leads to performance degradation
2. My current implementation requires to redo it for all relax models (no back-compatible)

[zxybazh]

LGTM, thanks for the quick update over OpenAI's update.