Sesame CSM in Tensorrt-llm

Hello,

I am interested in implementing a tensorrt-llm precompiled engine to serve the sesame csm tts model, and would like some guidance on ideas for how to implement and serve it. This model consists of two llama 3 models: a 1B backbone, and a 100M decoder. To generate a frames of audio from text input, it does the following:

1. Tokenizes text w Llama tokenizer, embeds the text tokens with a text token embedder
2. Passes text embeds to backbone. Backbone produces 1 audio token.
3. Backbone passes its audio token and its last hidden state to the decoder
4. The decoder then autoregressively generates 31 audio tokens (each token uses a different LM head)
5. These 32 tokens are then embedded (there is a separate embedding table for each token), the embeddings are summed into a single embedding, and this embedding is passed back into the backbone.

I want to be able to stream these blocks of 32 output tokens. I am confident in being able to build and compile the engine using the engine builder python api, but I am not positive what the best way to serve it is. I have the following two ideas:

The first is following the same style as the HF implementation. They do the following:

1. Implement the backbone as a LlamaModel with a custom embedding layer to handle both text and audio tokens.
2. Implement the decoder as a LlamaModelForCausalLM with a custom embedding layer and custom logic for the LM heads.
3. Here is where it gets hairy and where my main question is. They add custom generation that first generates the backbone logits normally. Then, they override the sample method to sample the backbone token, and run the 31 forward passes on the decoder model.

Would something like this be possible with the PyExecutor in Tensorrt-llm? Would it be possible to define a custom sampler which does the same thing as the huggingface implementation? It would sample the backbone logits, then run the decoder as a DecoderModelForCausalLM in a GenerationSession session for 31 iterations, and finally return the 32 tokens. This is kind of similar to speculative to decoding, minus the verification part.

The other strategy would be to serve the two separate engines and connect them using BLS in Triton Inference Server. Is that reasonable?

Thanks for any guidance