[OPTIMIZATION] Optimizes the single_query_cached_kv_attention kernel

[naed90]

Instead of having each thread group fetch the query head (which causes 64x memory to be read), we have all threads in the block share the task of loading the query head. On the benchmark of running 1000 sequences through LLaMA13B on an A100 (80GB), this improves the throughput by 1.10x.

[naed90]

See #421 for a detailed description and analysis of this commit.

[zhyncs]

Hi @naed90

overall LGTM, just had one small nitpick and looks like some formatting issues to address

[naed90]

Hi @naed90

overall LGTM, just had one small nitpick and looks like some formatting issues to address

ty.
can't seem to find your review, can you send a link to it?

[naed90]

@WoosukKwon @zhuohan123 hey, what do you think?

[WoosukKwon]

Hey @naed90, thanks for submitting the PR and apologies for the late response. I was busy for the last few days. Will take a look your issue and PR today.

[naed90]

Hey @naed90, thanks for submitting the PR and apologies for the late response. I was busy for the last few days. Will take a look your issue and PR today.

@WoosukKwon bump :)

[zhuohan123]

Tested a bit on the latency side:

Before optimization

$ python benchmark_latency.py --model huggyllama/llama-13b --input-len 128 --output-len 128 --num-iters 20
Namespace(model='huggyllama/llama-13b', tokenizer=None, tensor_parallel_size=1, input_len=128, output_len=128, batch_size=8, n=1, use_beam_search=False, num_iters=20, trust_remote_code=False)
INFO 07-24 21:53:31 llm_engine.py:67] Initializing an LLM engine with config: model='huggyllama/llama-13b', tokenizer='huggyllama/llama-13b', tokenizer_mode=auto, trust_remote_code=False, dtype=t
orch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)
INFO 07-24 21:53:31 tokenizer.py:29] For some LLaMA-based models, initializing the fast tokenizer may take a long time. To eliminate the initialization time, consider using 'hf-internal-testing/l
lama-tokenizer' instead of the original tokenizer.
INFO 07-24 21:54:01 llm_engine.py:183] # GPU blocks: 899, # CPU blocks: 327
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)
Warming up...
Profiling iterations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [01:11<00:00,  3.56s/it]
Avg latency: 3.5580986022949217 seconds

After optimization

$ python benchmark_latency.py --model huggyllama/llama-13b --input-len 128 --output-len 128
--num-iters 20
Namespace(model='huggyllama/llama-13b', tokenizer=None, tensor_parallel_size=1, input_len=128, output_len=128, batch_size=8, n=1, use_beam_search=False, num_iters=20, trust_remote_code=False)    INFO 07-24 21:55:36 llm_engine.py:67] Initializing an LLM engine with config: model='huggyllama/llama-13b', tokenizer='huggyllama/llama-13b', tokenizer_mode=auto, trust_remote_code=False, dtype=t
orch.float16, use_dummy_weights=False, download_dir=None, use_np_weights=False, tensor_parallel_size=1, seed=0)                                                                                    INFO 07-24 21:55:36 tokenizer.py:29] For some LLaMA-based models, initializing the fast tokenizer may take a long time. To eliminate the initialization time, consider using 'hf-internal-testing/l
lama-tokenizer' instead of the original tokenizer.
INFO 07-24 21:56:08 llm_engine.py:183] # GPU blocks: 899, # CPU blocks: 327
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)
Warming up...
Profiling iterations: 100%|████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████████| 20/20 [01:09<00:00,  3.49s/it]
Avg latency: 3.4891188383102416 seconds

[zhuohan123]

Thank you for your contribution! Left some small comments. We should be able to merge this after the changes.

[zhuohan123]

This if seems redundant if we assume NUM_THREADS should is divisible by THREAD_GROUP_SIZE?

[naed90]

Replaced with an assert.

[zhuohan123]

LGTM! Thank you again for your hard work and detailed profiling!