[RFC]C++ samplers and decoders support in Python

In the current TensorRT-LLM implementation, the C++ samplers and decoders support a wide variety of options(beam search/top-K/top-P/temperature/...). And with the recently introduced PyTorch-based architecture, support for the Python-native framework flavor is coming. We intend to unify these two designs keeping the flexibility that Python offers, and the performance of the C++ / CUDA optimizations.

Here is an overview of the current plan:

Fix functional issues of TRTLLMDecoder:

• Introduce synchronization, enable test for TRTLLMDecoder, and various fixes introduced in PR 3106.
• Fix configurations that require sending / receiving data from ranks. Examples include disaggregated serving and trtllm-bench with high concurrency.

Unify APIs of TRTLLMDecoder and TorchDecoder:

• The samplers / decoders should follow a common API, with the same inputs / outputs formats. This API should be followed in the implementation of the model, such as model_engine.py. The API should be:

class Decoder(ABC):

    @abstractmethod
    def setup_decoder_step(self, requests):
        raise NotImplementedError

    @abstractmethod
    def decode_async(self, scheduled_requests: ScheduledRequests,
                     model_outputs):
        raise NotImplementedError

    @abstractmethod
    def update_requests(self,
                        scheduled_requests: ScheduledRequests,
                        host_tensors: OrderedDict[str, torch.tensor],
                        decoder_event: torch.cuda.Event or None = None):
        raise NotImplementedError

• setup_decoder_step: Sets up the decoder requests and sampling configs, and assigns batch slots to the requests.
• decode_async: Performs a decoding step asynchronously and prepares the output in the expected format. The output should be a tuple (host_output, dev_output, decoder_event):

host_tensors: OrderedDict containing:	
	"new_tokens_host": tensor [max nb of tokens in each step][max batch size][max beam width]
	"finished_sum_host": tensor [max batch size] - Contains the number of finished beams for each seq slot.
	"finish_reasons_host": tensor [max batch size * max beam width] - Finish reason for each beam.
	"seq_lens_host": tensor [max batch size * max beam width] - Current seq len.

dev_tensors: OrderedDict containing the following entries:
	"new_tokens", "finished_sum", "finish_reasons", "seq_lens" - Shape similar to host ones described above.

decoder_event: CUDA event recorded after device -> host copies of above tensors.

• update_requests: Updates the LlmRequest objects accordingly. host_tensors should be as above.

Implement glue logic in Python:

• Convert C++ algorithms in setup: GenerateRequestOptions, CreateNewDecoderRequests.
• Convert C++ algorithms in forward: HandleContextLogits, HandleGenerationLogits, MakeDecodingBatchInputOutput.
• Outstandingly, GptDecoderBatched will remain, invoked with PyBind.
• Transition seamlessly to new Python-based LlmRequest and options objects.

Optimize Python-only features in C++ / CUDA:

• Going forward, the intent is to provide features first developed in Python for speed of development, and in C++ / CUDA for better performance.
• Core building blocks of MTP, Eagle-3 which can benefit from GPU acceleration will be ported to native CUDA to improve performance.

Feel free to leave your comments about this plan.

Thanks，

The TensorRT-LLM Engineering Team