fixes after rebase

serve/mlc_serve/model/model_common.py

@@ -176,10 +176,23 @@ def sample_from_logits(
     return outputs
 
 
+def get_logprob_infos(
+        i: int,
+        logprob_infos: Optional[List[Optional[RawLogprobsInfo]]],
+    ) -> Optional[List[Optional[RawLogprobsInfo]]]:
+        if logprob_infos is None or logprob_infos[i] is None:
+            return None
+        return [logprob_infos[i]]
+
+
 def sample_loglikelihood_from_logits(
     logits: Union[tvm.nd.NDArray, torch.Tensor],
     sequence_ids: List[SequenceId],
 ) -> List[TextGenerationResult]:
+    # Convert to torch tensors if logits are in tvm ndarray
+    if isinstance(logits, tvm.nd.NDArray):
+        logits = torch.from_dlpack(logits)
+
     # TODO(vvchernov): cut prompt lengths from logits
     try:
         logprob_infos = loglikelihood_sample(logits)

serve/mlc_serve/model/sampler.py

@@ -546,8 +546,6 @@ def get_loglikelihood_logprob_infos(
 def loglikelihood_sample(
     logits: torch.Tensor,
 ) -> List[Optional[RawLogprobsInfo]]:
-    logits = torch.from_dlpack(logits)
-
     # TODO(vvchernov): we need token_ids from input
     # It is not neccessary top1
     res = torch.argmax(logits, -1)

serve/mlc_serve/model/tvm_model.py

@@ -347,7 +347,7 @@ def generate(
                 request_past_decode_tokens = [self.vocab_size]
             elif isinstance(request, LoglikelihoodRequest):
                 seq_id = get_prompt_sequence_id(request.request_id)
-                # TODO(vvchernov): it it needed?
+                # TODO(vvchernov): is it needed?
                 request_past_decode_tokens = [self.vocab_size]
             elif isinstance(request, DecodeRequest):
                 seq_id = request.sequence_id