feat: return logits in PyTorch flow (NVIDIA#3221)

Signed-off-by: Yuan Tong <13075180+tongyuantongyu@users.noreply.github.com>
Co-authored-by: QI JUN <22017000+QiJune@users.noreply.github.com>

Author: tongyuantongyu

tensorrt_llm/_torch/models/modeling_deepseekv3.py

@@ -1096,7 +1096,7 @@ def forward(
                 attn_metadata,
                 True,
             )
-            # get accepetd tokens and next draft tokens
+            # get accepted tokens and next draft tokens
             return self.mtp_worker(
                 input_ids=input_ids,
                 position_ids=position_ids,

tensorrt_llm/_torch/pyexecutor/decoder.py

@@ -19,7 +19,7 @@
                                                     SpeculativeDecodingMode)
 from tensorrt_llm.mapping import Mapping
 
-from .llm_request import LlmRequest, LlmRequestState
+from .llm_request import LlmRequest, LlmRequestState, LogProbs
 from .scheduler import ScheduledRequests
 
 
@@ -29,6 +29,9 @@ class DecoderState:
 
     logits: torch.Tensor = None
 
+    # Set when decode_async() has evaluated these to avoid computing again in update_requests()
+    log_probs: list[LogProbs] | None = None
+
     new_tensors_device: dict[str, torch.Tensor] = None
     new_tensors_host: dict[str, torch.Tensor] = None
 
@@ -66,9 +69,13 @@ def update_requests(self, decoder_state: DecoderState) -> None:
         assert (not scheduled_requests.generation_requests)
         for idx, request in enumerate(scheduled_requests.context_requests):
             request.state = LlmRequestState.GENERATION_COMPLETE
-            #NOTE: This is a hack: set finish reason manually and set the beam 0
+            # NOTE: This is a hack: set finish reason manually and set the beam 0
             request.set_finished_reason(FinishReason.LENGTH, 0)
-            request.context_logits = decoder_state.logits[idx]
+            logits = decoder_state.logits[idx]
+            if logits.ndim == 1:
+                # For BERT: Add vocab_size axis to be compatible with LogitsStorage.
+                logits = logits.unsqueeze(-1)
+            request.py_result.append_context_logits(logits)
 
 
 def top_k_sampling_batch(logits, top_k=50):
@@ -158,10 +165,7 @@ def decode_single_request(request: LlmRequest, logits):
             logits, request.sampling_config.top_k[0])
     else:
         next_tokens, log_probs = greedy_search_sampling_batch(logits)
-    # TODO: enable these lines when log_probs is needed
-    # request.log_probs_async = log_probs
-    # request.set_cum_log_prob(request.cum_log_probs[0] + log_probs[0].item(), 0)
-    return next_tokens
+    return next_tokens, log_probs
 
 
 class TorchDecoder(Decoder):
@@ -222,23 +226,51 @@ def update_requests(self, decoder_state: DecoderState) -> None:
             "new_tokens_host"].tolist()
         scheduled_requests = decoder_state.scheduled_requests
 
-        idx = 0
+        request_idx = 0
+        token_idx = 0
         beam_idx = 0
+
+        def advance_idx(num_tokens=1):
+            nonlocal request_idx, token_idx
+            request_idx += 1
+            token_idx += num_tokens
+
+        def handle_logits(request: LlmRequest, count=1):
+            if decoder_state.logits is None:
+                return
+            if not request.py_return_generation_logits and not request.py_return_log_probs:
+                return
+
+            current_slice = slice(token_idx, token_idx + count)
+            current_logits = decoder_state.logits[current_slice]
+
+            request.py_result.append_generation_logits(current_logits)
+
+            if not request.py_return_log_probs:
+                return
+
+            if decoder_state.log_probs:
+                log_probs = decoder_state.log_probs[request_idx]
+            else:
+                _, log_probs = greedy_search_sampling_batch(current_logits)
+            request.py_result.append_log_probs([log_probs.tolist()])
+
         for request in scheduled_requests.context_requests:
             if request.get_context_remaining_length() != 0:
-                idx += 1
+                advance_idx()
                 continue
 
             if request.state != LlmRequestState.GENERATION_COMPLETE:
-                new_token = new_tokens_list[idx]
+                new_token = new_tokens_list[token_idx]
                 num_tokens = request.add_new_token(new_token, beam_idx)
                 self._handle_stop_criteria(request, new_token, num_tokens,
                                            beam_idx)
+                handle_logits(request)
                 request.py_decoding_iter += 1
-            idx += 1
+            advance_idx()
 
         if hasattr(scheduled_requests, 'chunked_requests'):
-            idx += len(scheduled_requests.chunked_requests)
+            request_idx += len(scheduled_requests.chunked_requests)
 
         extend_requests = []
         generation_requests = []
@@ -249,41 +281,41 @@ def update_requests(self, decoder_state: DecoderState) -> None:
                 generation_requests.append(request)
 
         for request in extend_requests:
-            num_accepted = 0
             if request.state != LlmRequestState.GENERATION_COMPLETE:
-                new_token = new_tokens_list[idx]
+                new_token = new_tokens_list[token_idx]
                 num_tokens = request.add_new_token(new_token, beam_idx)
                 self._handle_stop_criteria(request, new_token, num_tokens,
                                            beam_idx)
-                request.py_decoding_iter += 1
 
                 # Accept draft tokens (if we have any) if and only if they match the new
                 # token exactly.
-                for i in range(len(request.py_draft_tokens)):
-                    draft_token = request.py_draft_tokens[i]
+                num_accepted = 0
+                for draft_token in request.py_draft_tokens:
                     if draft_token != new_token:
                         # Reject.
                         break
-
                     num_accepted += 1
-                    new_token = new_tokens_list[idx + i + 1]
+                    new_token = new_tokens_list[token_idx + num_accepted]
                     num_tokens = request.add_new_token(new_token, beam_idx)
 
                     if self._handle_stop_criteria(request, new_token,
                                                   num_tokens, beam_idx):
                         break
-            request.py_num_accepted_draft_tokens = num_accepted
-            request.py_rewind_len = request.py_draft_pages_allocated - num_accepted
-            idx += len(request.py_draft_tokens) + 1
+                handle_logits(request, num_accepted)
+                request.py_decoding_iter += 1
+                request.py_num_accepted_draft_tokens = num_accepted
+                request.py_rewind_len = request.py_draft_pages_allocated - num_accepted
+            advance_idx(len(request.py_draft_tokens) + 1)
 
         for request in generation_requests:
             if request.state != LlmRequestState.GENERATION_COMPLETE:
-                new_token = new_tokens_list[idx]
+                new_token = new_tokens_list[token_idx]
                 num_tokens = request.add_new_token(new_token, beam_idx)
                 self._handle_stop_criteria(request, new_token, num_tokens,
                                            beam_idx)
+                handle_logits(request)
                 request.py_decoding_iter += 1
-            idx += 1
+            advance_idx()
 
     def _mixed_decode(self, scheduled_requests: ScheduledRequests,
                       model_outputs) -> DecoderState:
@@ -292,23 +324,33 @@ def _mixed_decode(self, scheduled_requests: ScheduledRequests,
         state = DecoderState(
             scheduled_requests=scheduled_requests,
             logits=logits,
+            log_probs=[],
         )
 
         new_tokens_device_array = []
 
         idx = 0
 
         for request in scheduled_requests.context_requests:
+            assert not request.py_return_context_logits, "Return context logits not supported"
             token_logits = logits[idx:idx + 1, :]
-            new_token = decode_single_request(request, token_logits)
+            new_token, log_probs = decode_single_request(request, token_logits)
             new_tokens_device_array.append(new_token)
+            log_probs = [log_probs.tolist()
+                         ] if request.py_return_log_probs else None
+            state.log_probs.append(log_probs)  # Currently always beam_width=1
             idx += 1
 
         for request in scheduled_requests.generation_requests:
+            if request.state == LlmRequestState.GENERATION_COMPLETE:
+                continue
             assert request.py_draft_tokens is None, "Speculative decoding not supported in SeparateDecoder."
             token_logits = logits[idx:idx + 1, :]
-            new_token = decode_single_request(request, token_logits)
+            new_token, log_probs = decode_single_request(request, token_logits)
             new_tokens_device_array.append(new_token)
+            log_probs = [log_probs.tolist()
+                         ] if request.py_return_log_probs else None
+            state.log_probs.append(log_probs)  # Currently always beam_width=1
             idx += 1
 
         new_tokens_device = torch.cat(new_tokens_device_array)
@@ -351,6 +393,13 @@ def update_one_request(self, request: LlmRequest,
         output_token_idx = request.output_token_idx
         new_token = new_tokens_list[output_token_idx]
         num_tokens = request.add_new_token(new_token, beam_idx)
+
+        current_logits = logits[output_token_idx].unsqueeze(0)
+        request.py_result.append_generation_logits(current_logits)
+        if request.py_return_log_probs:
+            _, log_probs = greedy_search_sampling_batch(current_logits)
+            request.py_result.append_log_probs([log_probs.tolist()])
+
         self._handle_stop_criteria(request, new_token, num_tokens, beam_idx)
         if request.state != LlmRequestState.GENERATION_COMPLETE:
             request.py_decoding_iter += 1