MTP and derivatives: Align sample state with trtllm sampler sample state (#5675)

This PR moves MTPSampler and derivatives to use the universal seq_slot indexing for sampling.
This is the last piece of the puzzle: After this, all of the samplers will use this format.
See: 6ee94c7
Signed-off-by: Netanel Haber <[email protected]>

Author: netanel-haber

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -25,6 +25,7 @@
 ExecutorRequest = tllm_executor.Request
 ExecutorResponse = tllm_executor.Response
 ExecutorSamplingConfig = tllm_executor.SamplingConfig
+FinishReason = tllm_executor.FinishReason
 
 REQUEST_TYPE_MAPPING = {
     tllm_executor.RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION:
@@ -319,6 +320,11 @@ def create_response(
     def is_dummy(self):
         return self.is_attention_dp_dummy or self.is_cuda_graph_dummy or self.is_dummy_request
 
+    def finish_by(self, reason: FinishReason, beam: int) -> None:
+        """CPP finish by reason does not support beam_width > 1"""
+        self.state = LlmRequestState.GENERATION_COMPLETE
+        self.set_finished_reason(reason, beam)
+
 
 def convert_wordlist(word_list) -> List[List[int]]:
     """Converts a wordlist from format:

tensorrt_llm/_torch/pyexecutor/model_engine.py

@@ -21,7 +21,6 @@
 import tqdm
 
 import tensorrt_llm.bindings.internal.userbuffers as ub
-from tensorrt_llm._torch.pyexecutor.llm_request import LlmRequest
 from tensorrt_llm._torch.pyexecutor.sampler import SampleStateTensors
 from tensorrt_llm._torch.speculative.mtp import SampleStateTensorsMTP
 from tensorrt_llm._utils import (is_trace_enabled, local_mpi_rank,
@@ -1174,16 +1173,6 @@ def _prepare_tp_inputs(
         draft_lens = []
         mrope_config = defaultdict(list)
 
-        mtp_batch_idx = 0  # Temporary: MTP (and Eagle3OneModel) remain the only samplers to index new_tokens serially
-
-        def py_batch_idx(request: LlmRequest) -> int:
-            if not self.without_logits:
-                return request.seq_slot
-            nonlocal mtp_batch_idx
-            batch_idx = mtp_batch_idx
-            mtp_batch_idx += 1
-            return batch_idx
-
         for request in scheduled_requests.context_requests:
             request_ids.append(request.py_request_id)
             all_prompt_tokens = request.get_tokens(0)
@@ -1213,7 +1202,7 @@ def py_batch_idx(request: LlmRequest) -> int:
                 ) if mrope_rotary_cos_sin.device == 'cpu' else mrope_rotary_cos_sin
                 mrope_config['mrope_rotary_cos_sin'].append(
                     mrope_rotary_cos_sin.to('cuda', non_blocking=True))
-            request.py_batch_idx = py_batch_idx(request)
+            request.py_batch_idx = request.seq_slot
 
         num_ctx_requests = len(scheduled_requests.context_requests)
         num_ctx_tokens = len(input_ids)
@@ -1295,11 +1284,11 @@ def py_batch_idx(request: LlmRequest) -> int:
                 num_cached_tokens_per_seq.append(past_seen_token_num)
                 request_ids.append(request.py_request_id)
                 # update batch index
-                request.py_batch_idx = py_batch_idx(request)
+                request.py_batch_idx = request.seq_slot
             else:
                 # update batch index
                 previous_batch_idx = request.py_batch_idx
-                request.py_batch_idx = py_batch_idx(request)
+                request.py_batch_idx = request.seq_slot
                 # inputs
                 # overlap scheduler can only support the speculative decoding
                 # methods with a fixed number of draft tokens
@@ -1350,7 +1339,7 @@ def py_batch_idx(request: LlmRequest) -> int:
             prompt_lengths.append(request.py_prompt_len)
             draft_lens.append(0)
 
-            request.py_batch_idx = py_batch_idx(request)
+            request.py_batch_idx = request.seq_slot
 
         previous_batch_len = len(previous_batch_indices)
 
@@ -1387,7 +1376,8 @@ def previous_seq_slots_device():
                 # previous input ids
                 previous_batch_tokens = previous_batch_len * (
                     1 + self.max_draft_len)
-                new_tokens = new_tokens_device[previous_slots, :].flatten()
+                new_tokens = new_tokens_device.transpose(
+                    0, 1)[previous_slots, :].flatten()
                 self.input_ids_cuda[num_tokens:num_tokens +
                                     previous_batch_tokens].copy_(
                                         new_tokens, non_blocking=True)

tensorrt_llm/_torch/pyexecutor/sampler.py

@@ -199,6 +199,10 @@ def add_token(request: LlmRequest,
     return new_token
 
 
+def int_tensor(shape: tuple[int, ...], device: str = 'cuda') -> torch.Tensor:
+    return torch.empty(shape, dtype=torch.int, device=device)
+
+
 class TorchSampler(Sampler):
     BEAM = 0
     MAX_BEAM_WIDTH = BEAM + 1
@@ -208,6 +212,9 @@ class Store:
         new_tokens: torch.Tensor
         """Shape: See cpp DecoderState.getAllNewTokens()"""
 
+    def create_store(self) -> Store:
+        return self.Store(new_tokens=int_tensor(self.NEW_TOKENS_SHAPE))
+
     @dataclass(frozen=True, kw_only=True)
     class Args:
         max_seq_len: int
@@ -223,18 +230,16 @@ def __init__(self, args: Args):
         assert args.max_beam_width == self.MAX_BEAM_WIDTH, "TorchSampler only supports beam_width = 1"
         self.num_seq_slots = args.max_num_sequences
 
+        self.NEW_TOKENS_SHAPE = (self.max_tokens, self.num_seq_slots,
+                                 self.MAX_BEAM_WIDTH)
         # AutoDeploy build creates the sampler in inference mode,
         # which would disallow in-place mutating of new_tokens.
         # So, we temporarily exit inference mode.
         with torch.inference_mode(False):
-            new_tokens = torch.zeros(
-                (self.max_tokens, self.num_seq_slots, self.MAX_BEAM_WIDTH),
-                dtype=torch.int,
-                device='cuda')
-            self.store = self.Store(new_tokens=new_tokens)
-
-    def _meet_max_token_stop_criteria(self, request: LlmRequest,
-                                      num_tokens: int):
+            self.store = self.create_store()
+
+    def _meet_max_token_stop_criteria(self, request: LlmRequest):
+        num_tokens = request.get_num_tokens(self.BEAM)
         return (num_tokens - request.py_orig_prompt_len
                 >= request.py_max_new_tokens) or (num_tokens
                                                   >= self.max_seq_len)
@@ -258,21 +263,20 @@ def _meet_stop_token_criteria(request: LlmRequest):
                     return True
         return False
 
-    def _handle_stop_criteria(self, request: LlmRequest, new_token: int, *,
-                              beam: int) -> bool:
+    def _handle_stop_criteria(self, request: LlmRequest,
+                              new_token: int) -> bool:
         """Handle stop criteria and set appropriate finish reasons and state.
         Returns True if generation should stop."""
         if new_token == request.py_end_id:
-            request.finish_by_reason(FinishReason.END_ID)
+            request.finish_by(FinishReason.END_ID, self.BEAM)
             return True
 
-        num_tokens = request.get_num_tokens(beam)
-        if self._meet_max_token_stop_criteria(request, num_tokens):
-            request.finish_by_reason(FinishReason.LENGTH)
+        if self._meet_max_token_stop_criteria(request):
+            request.finish_by(FinishReason.LENGTH, self.BEAM)
             return True
 
         if self._meet_stop_token_criteria(request):
-            request.finish_by_reason(FinishReason.STOP_WORDS)
+            request.finish_by(FinishReason.STOP_WORDS, self.BEAM)
             return True
 
         return False
@@ -307,7 +311,7 @@ def process_draft_tokens(self, request: LlmRequest,
                                   new_tokens,
                                   beam=self.BEAM,
                                   step=num_accepted)
-            if self._handle_stop_criteria(request, new_token, beam=self.BEAM):
+            if self._handle_stop_criteria(request, new_token):
                 break
         return num_accepted
 
@@ -321,15 +325,15 @@ def update_requests(self, state: SampleState) -> None:
             if req.state == LlmRequestState.GENERATION_COMPLETE or req.context_remaining_length != 0:
                 continue
             new_token = add_token(req, new_tokens, beam=self.BEAM)
-            stop = self._handle_stop_criteria(req, new_token, beam=self.BEAM)
+            self._handle_stop_criteria(req, new_token)
             self.handle_logits(req, state, beam=self.BEAM, count=1)
             req.py_decoding_iter += 1
 
         for req in state.scheduled_requests.generation_requests:
             if req.state == LlmRequestState.GENERATION_COMPLETE:
                 continue
             new_token = add_token(req, new_tokens, beam=self.BEAM)
-            stop = self._handle_stop_criteria(req, new_token, beam=self.BEAM)
+            stop = self._handle_stop_criteria(req, new_token)
             processed = 1
             if not stop and len(req.py_draft_tokens) > 0:
                 num_accepted = self.process_draft_tokens(

tensorrt_llm/_torch/speculative/eagle3.py

@@ -472,13 +472,15 @@ def draft_decoder(
                 Draft token ids. Flattened.
         '''
 
-        draft_tokens = torch.argmax(logits, dim=-1).type(torch.int32)
+        draft_tokens = torch.argmax(logits, dim=-1)
 
         # Apply d2t (offsets between draft model dictionary and main model dictionary).
         if hasattr(draft_model.model,
                    "d2t") and draft_model.model.d2t is not None:
             draft_tokens = draft_model.model.d2t[draft_tokens] + draft_tokens
 
+        draft_tokens = draft_tokens.type(torch.int32)
+
         return draft_tokens
 
     def prepare_1st_drafter_inputs(

tensorrt_llm/_torch/speculative/mtp.py

@@ -4,12 +4,11 @@
 import torch
 from torch import nn
 
-from tensorrt_llm.bindings.executor import FinishReason
-
 from ..attention_backend import AttentionMetadata
 from ..pyexecutor.llm_request import LlmRequest, LlmRequestState
 from ..pyexecutor.resource_manager import BaseResourceManager, SlotManager
-from ..pyexecutor.sampler import SampleState, SampleStateTensors, TorchSampler
+from ..pyexecutor.sampler import (SampleState, SampleStateTensors, TorchSampler,
+                                  add_token, int_tensor)
 from ..pyexecutor.scheduler import ScheduledRequests
 from .interface import SpecConfig, SpecMetadata, SpeculativeDecodingMode
 
@@ -249,92 +248,96 @@ class MTPSampler(TorchSampler):
     SampleState = SampleStateMTP
 
     def __init__(self, args: TorchSampler.Args, *, nextn: int):
-        super().__init__(args)
         self.mapping = None
         self.draft_len = nextn
+        super().__init__(args)
 
-    def _draft_meet_max_token_stop_criteria(self, request: LlmRequest,
-                                            num_tokens: int, beam_idx: int):
-        if self._meet_max_token_stop_criteria(request, num_tokens):
-            request.state = LlmRequestState.GENERATION_COMPLETE
-            request.set_finished_reason(FinishReason.LENGTH, beam_idx)
+    @dataclass(frozen=True, kw_only=True)
+    class Store(TorchSampler.Store):
+        next_new_tokens: torch.Tensor
+        next_draft_tokens: torch.Tensor
+        new_tokens_lens: torch.Tensor
+
+    def create_store(self) -> Store:
+        num_tokens, seq_slots, _ = self.NEW_TOKENS_SHAPE
+        draft_len = num_tokens - 1
+        assert draft_len == self.draft_len
+        return self.Store(
+            new_tokens=int_tensor(self.NEW_TOKENS_SHAPE),
+            next_new_tokens=int_tensor(self.NEW_TOKENS_SHAPE),
+            next_draft_tokens=int_tensor((seq_slots, draft_len)),
+            new_tokens_lens=int_tensor((seq_slots, )),
+        )
+
+    def _request_common_handling(self, request: LlmRequest,
+                                 next_draft_tokens: list[list[int]]):
+        assert not request.py_return_context_logits, "return_context_logits not implemented for MTPSampler"
+        assert not request.py_return_generation_logits, "return_generation_logits not implemented for MTPSampler"
+        assert not request.py_return_log_probs, "return_log_probs not implemented for MTPSampler"
+        request.py_draft_tokens = next_draft_tokens[request.seq_slot]
+        request.py_decoding_iter += 1
 
     def update_requests(self, state: SampleStateMTP) -> None:
         assert isinstance(state, SampleStateMTP)
 
         state.sampler_event.synchronize()
-        new_tokens_list = state.host.new_tokens.tolist()
-        new_tokens_lens_list = state.host.new_tokens_lens.tolist()
+        new_tokens = state.host.new_tokens
+        new_tokens_lens = state.host.new_tokens_lens
         next_draft_tokens_list = state.host.next_draft_tokens.tolist()
-
-        idx = 0
-        beam_idx = 0
-        for request in state.scheduled_requests.context_requests:
-            assert not request.py_return_context_logits, "return_context_logits not implemented for MTPSampler"
-            assert not request.py_return_generation_logits, "return_generation_logits not implemented for MTPSampler"
-            assert not request.py_return_log_probs, "return_log_probs not implemented for MTPSampler"
-            if request.context_remaining_length != 0:
-                idx += 1
+        beam_idx = self.BEAM
+        for req in state.scheduled_requests.context_requests:
+            if req.state == LlmRequestState.GENERATION_COMPLETE or req.context_remaining_length != 0:
                 continue
+            new_token = add_token(req, new_tokens, beam=beam_idx)
+            self._handle_stop_criteria(req, new_token)
+            self._request_common_handling(req, next_draft_tokens_list)
 
-            if request.state != LlmRequestState.GENERATION_COMPLETE:
-                new_token = new_tokens_list[idx][0]
-                num_tokens = request.add_new_token(new_token, beam_idx)
-                should_stop = self._handle_stop_criteria(request,
-                                                         new_token,
-                                                         beam=beam_idx)
-                if self._draft_meet_max_token_stop_criteria(
-                        request, num_tokens, beam_idx):
-                    should_stop = True
-                request.py_draft_tokens = next_draft_tokens_list[idx]
-                request.py_decoding_iter += 1
-            idx += 1
-
-        for request in state.scheduled_requests.generation_requests:
-            assert not request.py_return_context_logits, "return_context_logits not implemented for MTPSampler"
-            assert not request.py_return_generation_logits, "return_generation_logits not implemented for MTPSampler"
-            assert not request.py_return_log_probs, "return_log_probs not implemented for MTPSampler"
-            if request.state != LlmRequestState.GENERATION_COMPLETE:
-                new_tokens = new_tokens_list[idx]
-                num_new_tokens = new_tokens_lens_list[idx]
-                should_stop = False
-                for i in range(num_new_tokens):
-                    new_token = new_tokens[i]
-                    num_tokens = request.add_new_token(new_token, beam_idx)
-                    should_stop = self._handle_stop_criteria(request,
-                                                             new_token,
-                                                             beam=beam_idx)
-                    if should_stop:
-                        break
-                if self._draft_meet_max_token_stop_criteria(
-                        request, num_tokens, beam_idx):
-                    should_stop = True
-                request.py_draft_tokens = next_draft_tokens_list[idx]
-                request.py_rewind_len = self.draft_len - (num_new_tokens - 1)
-                request.py_decoding_iter += 1
-            idx += 1
+        for req in state.scheduled_requests.generation_requests:
+            if req.state == LlmRequestState.GENERATION_COMPLETE:
+                continue
+            num_new_tokens = new_tokens_lens[req.seq_slot]
+            for i in range(num_new_tokens):
+                new_token = add_token(req, new_tokens, beam=beam_idx, step=i)
+                if self._handle_stop_criteria(req, new_token):
+                    break
+            req.py_rewind_len = self.draft_len - (num_new_tokens - 1)
+            self._request_common_handling(req, next_draft_tokens_list)
 
     def sample_async(self, scheduled_requests: ScheduledRequests,
-                     model_outputs) -> SampleStateMTP:
-        # new_tokens_device: all of the accepted tokens, device tensor
-        # new_tokens_lens_device: the accepted lengths, device tensor
-        # next_draft_tokens_device: predicted draft tokens, device tensor
-        # next_new_tokens_device: input tokens for the next iteration, device tensor
-        new_tokens_device = model_outputs['new_tokens']
-        new_tokens_lens_device = model_outputs['new_tokens_lens']
-        next_draft_tokens_device = model_outputs['next_draft_tokens']
-        next_new_tokens_device = model_outputs['next_new_tokens']
+                     outputs: dict[str, torch.Tensor]) -> SampleStateMTP:
+        # new_tokens_device: accepted tokens, device tensor, shape: batch_size, nextn + 1
+        # new_tokens_lens_device: accepted lengths, device tensor, shape: batch_size
+        # next_draft_tokens_device: predicted draft tokens, device tensor, shape: batch_size, nextn
+        # next_new_tokens_device: input tokens for the next iteration, device tensor, shape: batch_size, nextn + 1
+
+        requests = scheduled_requests.all_requests()
+        slots = torch.as_tensor([r.seq_slot for r in requests])
+        slots = slots.to(device="cuda", non_blocking=True)
+
+        o_new_tokens = outputs['new_tokens'][:len(requests)]
+        o_new_tokens_lens = outputs['new_tokens_lens'][:len(requests)]
+        o_next_draft_tokens = outputs['next_draft_tokens'][:len(requests)]
+        o_next_new_tokens = outputs['next_new_tokens'][:len(requests)]
+
+        new_tokens = self.store.new_tokens
+        next_new_tokens = self.store.next_new_tokens
+        new_tokens_lens = self.store.new_tokens_lens
+        next_draft_tokens = self.store.next_draft_tokens
+
+        new_tokens.squeeze(-1).T.index_copy_(0, slots, o_new_tokens)
+        next_new_tokens.squeeze(-1).T.index_copy_(0, slots, o_next_new_tokens)
+        new_tokens_lens.index_copy_(0, slots, o_new_tokens_lens)
+        next_draft_tokens.index_copy_(0, slots, o_next_draft_tokens)
 
         device = SampleStateTensorsMTP(
-            new_tokens=next_new_tokens_device,
-            new_tokens_lens=new_tokens_lens_device,
-            next_draft_tokens=next_draft_tokens_device,
+            new_tokens=next_new_tokens,
+            new_tokens_lens=new_tokens_lens,
+            next_draft_tokens=next_draft_tokens,
         )
         host = SampleStateTensorsMTP(
-            new_tokens=new_tokens_device.to('cpu', non_blocking=True),
-            new_tokens_lens=new_tokens_lens_device.to('cpu', non_blocking=True),
-            next_draft_tokens=next_draft_tokens_device.to('cpu',
-                                                          non_blocking=True),
+            new_tokens=new_tokens.to('cpu', non_blocking=True),
+            new_tokens_lens=new_tokens_lens.to('cpu', non_blocking=True),
+            next_draft_tokens=next_draft_tokens.to('cpu', non_blocking=True),
         )
         sampler_event = torch.cuda.Event()
         sampler_event.record()