Optimized Parler TTS

This repository aims to optimize inference of the Parler TTS family of models and implements Parler TTS from scratch borrowing code from

1. nanoGPT ¹
2. gpt fast ²
3. Parler TTS ³

A rough architecture of Parler TTS is as follows:

1. Text encoder: Parler TTS uses Flan T5 encoder⁴ for text encoding. This is a ~340M parameters model (encoder only). A single pass is sufficient for encoding, so the run time is extremely fast. This is used to encode the description of the prompt.
   1. This can be made faster by integrating flash attention into the attention layer.
   2. This is not however the slowest part in the complete pipeline taking <300ms for encoding 128 tokens.
   3. For v1, I am not speeding up this model; instead, I am just loading the model in bfloat16.
2. (Decoder) MusicGen autoregressive decoder⁵: This is the main decoder which produces audio tokens. This is ~470M params and decodes audio tokens autoregressively. This is the primary bottleneck and hence the main focus of this work.
3. Descript DAC Audio tokenizer⁶: This is the final decoder that takes in the generated audio codes from the above model, and generates audio from that. This is ~75M params and an extremely lightweight model. This is not the bottleneck for audio generation and hence not a focus of this work.

The decoder model takes 2 inputs:

1. Prompt - The text that needs to be converted to speech. These are encoded with a projector layer and fed directly to the model.
2. Description - A description of how the prompt should be spoken. These are encoded with Flat T5 and used for cross-attention in the attention layers.

At every decoding step, the decoder generates a token from each of the codebooks of DAC parallelly following a delayed pattern resulting in the generation of 9 tokens at every step:

Source: Simple and Controllable Music Generation ⁴

Performance Analysis

Architectural quirks

The decoder takes prompts and descriptions as embeddings. This leads to a complicated pipeline that requires maintaining an additional model for each forward pass. Instead, if the decoder can be trained on text and audio tokens directly, it would avoid the need for an external text encoder. Another possibility to make the model inference more streamlined is by converting the description to tags and using them as tokens while decoding. This would avoid the cross-attention with description embeddings and make the architecture much simpler. This would also avoid costly cross-attention if the description text is too big.

Speeding up Parler TTS decoder

The following analysis was done on parler-tts/parler-tts-mini-v1 running on Nvidia RTX 3080 Ti.

Using the official repository, I get the following numbers:

(Average of 10 runs)

We can see that as the number of input tokens increases (and the description tokens), the throughput decreases due to cross-attention on the complete description. However, as the output tokens increase, the throughput does not go down drastically. In the real world, if the description is too long, the throughput and time to the first token will suffer greatly.

Analysis and improvements

The FLOP requirements for the mini v1 model assuming fp16 cores and an input sequence length of 64 are ~36 GFLops. The memory requirement for the same model is 950 MB assuming fp16 weights. For 3080 Ti (912 GB/s memory bandwidth and 34 TFlops fp16 compute), I get the following latencies: $L_{memory_bandwidth} = 1 * 10^{-3}s$ and $L_{compute} = 1.1 * 10^{-3}s$. This equates to ~800 toks/s (I am assuming a bit here, discounting for KV cache movement). Looking at the FLOPs requirement and the memory bandwidth requirement, the model does not seem to be either memory-bound or compute-bound on the current generation speeds.

On looking at the profile of the decoder, you can see lots of gaps in the GPU stream:

This is because the GPU is extremely fast at executing the kernels and is kept waiting for the CPU to send it to the next kernel. This leads to the model being overhead-bound. To reduce the overhead bound, we can capture the complete execution of the model as a single graph and launch all the kernels at once as CUDA graphs.

To implement the CUDA graphs, I implemented the following:

1. Static caching for both Cross Attention and Causal Attention. This is necessary to keep all the tensor allocations static so that we can compile the model as a single CUDA graph
2. Sampling without explicit synchronization call between GPU and CPU

After compilation, I get the following profile of the decoder:

We can now see that all the kernels are being executed one after the other which speeds up the overall execution and drastically increases the throughput. You can also see how multiple kernels are being executed from a single 'cudaGraphLaunch` instruction from the CPU.

However, one drawback of compiling the model with static caches is that we have to preallocate the memory for the KV caches which can be huge depending on the batch size and the embedding dimension. This also leads to the attention kernel taking more time to execute than eager mode because now it has to transfer the complete KV cache for every forward pass.

In the official implementation, it takes ~13 microseconds for the attention kernel.

VS

In my custom implementation, it took ~95 microseconds for the attention kernel.

To reduce the load on the memory bandwidth, I tried quantizing the model as well. I quantized the model to int8 to see if the attention kernels would take less time or not. Surprisingly, it took a similar amount of time and the overall execution time increased. I suspect this is due to the computation overhead that is associated with converting the weights back to bf16 for computation.

Other than this, I also implemented a few other optimizations:

1. For the language model head, only do the forward pass on the last logit (since that is the only new token we would decode)
2. There are 9 separate heads for the language model, which in the original implementation are used one by one (ref). I merged all the 9 heads into a single weight matrix and did the forward pass only once. This had a minimal impact on the throughput.
3. Similarly, there are 9 separate embedding lookups (ref). I merged them into 1 lookup instead of 9 separate ones. This did not have any impact on the throughput.
4. Increased the size of the embedding matrix to be a multiple of 8 (it was 1089 earlier). This can help in better distribution of work in an SM (ref)

Overall, after implementing all of the above, I get the following improvements.

(Average of 10 runs)

I have also implemented batching, but I am unable to verify the output of the audio for the complete batch.

Finally, on a batch size of 1, these are the results:

Keeping output tokens fixed, I achieved an improvement of 10x in end-to-end latency while decoding 1024 tokens (11s of audio) with 1024 tokens as input.

Keeping input tokens fixed, I achieved an improvement of 1.3x in end-to-end latency while decoding 2048 tokens (23s of audio) with 64 tokens as input.

I achieve an RTFx of 1.3 on the extreme cases in both the scenarios whereas, HuggingFace's implementation gets an RTFx of < 0.4 for long input sequences and < 1 for long output tokens.

Streaming

I have replicated the official implementation of the streamer that can decode the audio tokens as they come and stream the audio. The decoder puts the audio tokens to a queue and the audio tokenizer decodes all the collected tokens into audio.

I have currently not integrated this streamer with my custom implementation of ParlerTTS.

Future work

1. The current implementation implements a static cache that can put an unnecessary load on the memory bandwidth. This can be avoided by implementing PagedAttention kernels that can reduce the cost of cache transfer.
2. Sampling is currently very inefficient which sends multiple kernels (>10) for temperature scaling, Top K sampling, and Softmax. This leads to multiple reads and writes from global memory and can slow down the overall generation speed. This can be improved by fusing these kernels into one. Even the torch compiler fails to fuse these kernels into one.

References

Footnotes

1. nano GPT ↩

2. gpt fast ↩

3. Official implementation ↩

4. Scaling Instruction-Finetuned Language Models ↩ ↩²

5. Simple and Controllable Music Generation ↩

6. Descript audio codec ↩