Mask cache Performance Optimization for vllm

[paul-grundmann]

Problem

The current implementation allocates a mask for every token during generation, which significantly impacts performance.

Proposed Solution

To improve the performance, we can cache the mask on the device, as it depends on the allowed tokens from the FSM. Additionally, limiting the input to the hash function to the first 2k tokens results in a notable speedup.

Discussion

While using only the first 2k tokens for the hash may introduce potential cache collisions, the likelihood of such collisions is very low.

TODO

• Provide measurements of the performance impact

[rlouf]

My worry with limiting the number of tokens for the hash is that we might return identical masks for different lists of tokens; it is not excluded that first 2048 allowed tokens would be the same but next ones different.

[paul-grundmann]

I've made additional optimizations and determined that caching based solely on the state id is sufficient. This eliminates the need for the sometimes expensive hashing of allowed input tokens. As a result, on our A100 and H100 hardware, generation performance is nearly as fast as it is without a logits_processor

[rlouf]

Could you run pre-commit locally?

[lapp0]

Could you please rebase onto main? Your diff is currently the entire repo.

[paul-grundmann]

Rebase is done. Don't know what happened there.

Benchmark

I also ran some benchmarks on 128 samples using an H100 PCIe GPU using LLaMA3 8B:

Without Cache:
128/128 [03:46<00:00,  1.77s/it, Generation Speed: 81.87 toks/s]

Cold Cache:
128/128 [00:41<00:00,  3.11it/s, Generation Speed: 451.46 toks/s]

Warm Cache:
128/128 [00:13<00:00,  9.74it/s, Generation Speed: 1412.09 toks/s]

In addition here are the plots of the GPU utilization during the benchmark:

No cache

Cold cache

Warm cache

As depicted in the plots: Without caching the masks, the gpu utilization remains very low. A fully warmed up cache results in an improvement of over 17x and an almost fully used GPU (in terms of GPU utilization and power consumption). I would estimate that the difference could be even higher for larger amounts of samples.

[rlouf]

Looks good to me, thank you for contributing!