Motivation

We suspect the concatenations of the key/value buffer at each generation step to be expensive.

FasterTransformer does preallocate the kv cache. Preallocating may also help torch.compile, also I don't quite get why (there are still dynamic shapes in the model itself, so why care about the model I/O?).

Results

Some optimizations may still be missing. On commit 1acc7e4ea234d7143700234e51877acc9859f9ed

Todos

• past_key_values buffer should be initialized in the background, instead of requiring the user to initialize it himself and requiring to pass generate(**inputs, past_key_values=past_key_values).
• Current implementation is likely to break with accelerate with naive pipeline parallelism, as the buffer is currently initialized on a single device
• Preallocated kv cache still does not help with small models / batch size.
• Support an iterative buffer, e.g. that auto-extends each 512 tokens, instead of initializing a buffer of size max_new_tokens. This may help reducing memory usage (and speed? probably not).
• Implement tests
• Should this be in optimum or transformers?
• Support all (is it possible?) decoding strategies, instead of currently only greedy_search
• Have it work with cross-attention

• Experiment with torch.compile (find torchdynamo graph breaks and remove them if possible?)
• Experiment with iteratively allocated KV cache, see this suggestion, this would avoid aten::slice calls
• Can we avoid aten::copy_ calls and aten::slice calls?
• Test on CPU
• Support cross-attention
• Support encoder-decoder architectures
• Support preallocated attention_mask
• Support preallocated token_type_ids (llama does not use it though)
• Would a single qkv_proj help instead of the current q_proj separate from kv_proj? [yes, helps slightly]
• Make sure the implementation is valid not on the argmax but on the logits directly
• Investigate whether stopping to inherit from nn.Module may help, see this __getattr__ overhead