Optimize MQA Kernel

[zhuohan123]

This PR implements the MQA paged attention kernel and modifies the GPT Bigcode model to utilize the optimized MQA kernel.

TODO: Check performance gain.

[nivibilla]

How hard is it to port MQA to LLaMA models?

[WoosukKwon]

Awesome! Thanks for the PR. Left some comments.

[WoosukKwon]

Style nit:

Suggested change

	assert self.num_heads % self.num_kv_heads == 0
	self.head_mapping = torch.repeat_interleave(
	torch.arange(self.num_kv_heads, dtype=torch.int32, device="cuda"),
	num_heads // self.num_kv_heads)
	assert self.num_heads % self.num_kv_heads == 0
	self.num_queries_per_kv = self.num_heads // self.num_kv_heads
	self.head_mapping = torch.repeat_interleave(
	torch.arange(self.num_kv_heads, dtype=torch.int32, device="cuda"),
	self.num_queries_per_kv)

self.num_queries_per_kv can be also used in L97 and L100.

[zhuohan123]

Starcoder latency after this PR on 1 GCP A100:

$ python benchmark_latency.py --model bigcode/starcoder --batch-size 1 --input-len 128 --output-len 128 --num-iters 1
Namespace(model='bigcode/starcoder', tokenizer=None, tensor_parallel_size=1, input_len=128, output_len=128, batch_size=1, n=1, use_beam_search=False, num_iters=1, profile=False)
INFO 07-17 07:23:28 llm_engine.py:60] Initializing an LLM engine with config: model='bigcode/starcoder', tokenizer='bigcode/starcoder', tokenizer_mode=auto, trust_remote_code=False, dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
INFO 07-17 07:25:40 llm_engine.py:134] # GPU blocks: 20280, # CPU blocks: 13107
SamplingParams(n=1, best_of=1, presence_penalty=0.0, frequency_penalty=0.0, temperature=1.0, top_p=1.0, top_k=-1, use_beam_search=False, stop=[], ignore_eos=True, max_tokens=128, logprobs=None)
Avg latency: 3.6188764572143555 seconds

before this PR:

(baseline-branch) ubuntu@zhuohan-vllm-4-manual:~/nfs/cacheflow/base-branch/vllm/benchmarks$ python benchmark_latency.py --model bigcode/starcoder --batch-size 1 --input-len 128 --output-len 128
--num-iters 1
Namespace(model='bigcode/starcoder', tokenizer=None, tensor_parallel_size=1, input_len=128, output_len=128, batch_size=1, n=1, use_beam_search=False, num_iters=1)
INFO 07-17 07:28:23 llm_engine.py:60] Initializing an LLM engine with config: model='bigcode/starcoder', tokenizer='bigcode/starcoder', tokenizer_mode=auto, trust_remote_code=False, dtype=torch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
INFO 07-17 07:30:55 llm_engine.py:134] # GPU blocks: 49, # CPU blocks: 273
SamplingParams(n=1, best_of=1, presence_penalty=0.0, frequency_penalty=0.0, temperature=1.0, top_p=1.0, top_k=-1, use_beam_search=False, stop=[], ignore_eos=True, max_tokens=128, logprobs=None)
Avg latency: 4.571176528930664 seconds

Note that the throughput can be further boosted because of the extra cache blocks provided.

[zhyncs]

Hi @zhuohan123

May I ask, why implement the MQA kernel optimization in the single_query_cached_kv_attention_kernel instead of abstracting a kernel separately

[zhuohan123]

Hi @zhuohan123

May I ask, why implement the MQA kernel optimization in the single_query_cached_kv_attention_kernel instead of abstracting a kernel separately

Hi! MQA can be combined with some other attention modifications, including RoPE and Alibi embeddings. Having a separate class for MQA can introduce many new classes with these embedding variations.