LLM 인퍼런스 훑어보기 (3) - KV Cache (deep dive)

이전 포스트에서는 트랜스포머의 구조와 어텐션 메커니즘이 무엇인지 다루었다. 그리고 LLM 인퍼런스에서 KV Cache가 무엇이며, 왜 도입되어 사용되는지에 대해서 다뤘다. 지난 포스트는 아래 링크를 참조하자.

LLM 인퍼런스 훑어보기 (2) - KV Cache

 

이번 포스트에서는 다른 가속 방법을 다루려고 했었는데, KV Cache에 대해서 좀 더 deep dive 해보는 내용을 작성해보려고 한다. deep dive 내용은 llama 모델의 구현체 위주로 살펴볼 예정이며, 크게 huggingface transformers의 구현체와 vLLM 프로젝트의 구현체를 살펴본다. 

 

---

 

Huggingface Transformers에서의 KV Cache

huggingface transformers에서 제공하는 llama 모델의 causalLM 형태는 LlamaForCausalLM 객체로 제공하고 있다.

 

LlamaForCausalLM 객체는 멤버로 LlamaModel을 가지고 있고, 이 model의 forward call을 사용하여 forward 한다.

 

기본적으로 hugginface transformers의 PreTrainedModel 객체는 GenerationMixin이라는 Mixin 객체를 상속받는다. 이 GenerationMixin에는 generation과 관련된 메소드들이 포함되어 있다.

 

그 중 우리는 generate라는 함수를 살펴볼 것이다.

이 함수는 위와 같은 인자를 갖고 있다. 인자 하나하나에 대한 설명은 이 포스트에서 다룰 내용은 아니므로, 다음번에 관련해서 별도의 포스트를 작성하겠다. 본론으로 돌아가서 KV cache와 관련된 구현체만 살펴보자

 

decoder-only 모델의 경우 input_embeds로 입력이 들어오는 경우에는 cache를 무조건 사용해야만 한다. 왜냐하면 입력값에 대해서 첫 토큰을 생성하고 있는지 아닌지 확인할 수 없기 때문이다. 그 외에는 설정으로 들어온대로 사용한다.

 

이 내용은 크게 상관은 없지만, encoder_decoder 모델의 경우에는 사전에 미리 encoder에서 output을 추출해서 hidden_states와 attention 값들을 model_kwargs에 저장해둔다.

 

 

decoding 전략은 다른 부분은 배제하기 위해서, 가장 간단한 multinomial sampling을 위주로 살펴보자.

 

sampling에서 while loop를 타기전에 model_kwargs에 저장되어 있는 encoder_ouputs을 추출해서 준비한다.

 

많은 부분을 생략하긴 했지만..., while loop에서 진행되는 것들을 요약하자면 다음과 같다:

1. 우선, input_ids와 model_kwargs을 베이스로 prepare_inputs_for_generation 함수를 거쳐 generation을 위한 model_inputs를 생성한다.
2. 생성된 model_inputs는 model의 입력 값으로 주입되어 outputs를 생성해낸다.
3. 이렇게 생성된 outputs는 model_kwargs와 is_encoder_decoder를 인자로 갖는 _update_model_kwargs_for_generation 함수를 통해 새로운 model_kwargs로 갱신된다.

 

_update_model_kwargs_for_generation 함수를 살펴보면 더 자세한 내용을 파악할 수 있다.

_update_model_kwargs_for_generation 함수에서는 outputs에서 past를 추츨하여 past_key_values라는 kv cache를 갱신한다! 물론 token_type_ids, attention_mask, decoder_attention_mask 등을 갱신하는 작업도 수행한다.

 

그럼 이제 while loop에서 past_key_values로 캐시를 갱신하는 건 알겠는데, 이 것을 사용하는 부분을 보자. 사용하는 부분은 llama의 모델링 구현체를 살펴봐야 한다.

아래는 LlamaModel 객체의 forward 중 일부이다. forward에서 past_key_values 인자로 값이 전달된다면, past_key_values_length를 past_key_values[0][0].shape[2]로 갱신한다. llama에서 past_key_values[0][0]은 key_states를 의미하고, key_states는 [batch_size, num_key_value_heads, q_len, head_dim]이기 때문에 past_key_values_length=q_len이다. 즉 past_key_values_length는 말 그대로 past_key_values의 길이를 말한다.

 

decoding layer에 past_key_values가 어떻게 적용되는지를 확인해보자.

1. decoding layer의 각 층을 순회하는 for loop
2. past_key_value는 past_key_values에서 추출한 각 레이어에 주입할 cache값
3. past_key_value는 decoding layer에 인자로 전달되어 사용
4. next_decoder_cache는 다음 generation loop에 전달하기 위해 층마다 출력된 key & value states를 append
5. 최종적으로 BaseModelOutputWithPast 객체에 past_key_values값으로 next_decoder_cache값 전달

이렇게 ModelOutput으로 전달된 past_key_values는 generation loop에서 갱신되며 계속 사용된다.

 

next_decoder_cache에 저장되는 값은 LlamaAttention 레이어가 출력하는 past_key_value로 확인할 수 있는데, 다음과 같다:

• 어텐션 레이어 단에서 past_key_value를 주입받아서 key_states와 values_states를 past_key_values에서 추출하여 dim=2 concat하여 사용하게 된다(dim=2는 q_len).
• concat된 key_states와 values_states는 past_key_value로 tuple로 다시 묶여 반환된다.

 

이상으로 Huggingface Transformers에서 KV cache를 다루는 구현체를 알아봤다. kv cache는 past_key_values라는 이름으로 저장하고 있으며, model_kwargs라는 object에 저장되는 것을 확인할 수 있었다.

 

---

 

vLLM에서의 KV Cache

다음으로 vLLM에서의 KV Cache가 어떻게 구현되어 있는지 확인해보자. vLLM에서는 modelling 구현체에서 KV Cache의 타입이 정의되어 있다. 

 

위의 KVCache 자료형은 huggingface transformers에서처럼 (key, value)로 구성되어 있다.

 

vLLM에 구현된 LlamaForCauslaLM도 forward에서 kv_cache를 인자로 전달받는다.

 

vLLM은 KV Cache를 잘 관리하고 이러한 cache에서부터 pagedAttention을 활용하는게 주 목적인 프레임워크라서 cache를 처리하는 다양한 로직들이 포함되어 있다. 이러한 로직은 worker의 execute_model 함수에 추가되어 있다.

1. cache_engine에서 swap_in, swap_out, copy에 대한 처리가 필요한 경우를 cuda stream으로 이벤트를 발생시킨다.
   • swap_in: cpu_cache block들을 gpu_cache로 옮기는 것
   • swap_out: gpu_cache block들을 cpu_cache로 옮기는 것
   • copy: cpu와 gpu cache를 동기화하는 것
2. 이러한 이벤트를 cache_events라고 하며 이를 모델에 전달한다.
3. 실질적으로 모델의 kv_cache는 gpu_cache를 바라본다.

 

아래는 vLLM에서 구현한 pagedAttention 중 cache_event를 기다리는 내용이다. 즉 layer별로 offload 등과 같은 cache에 대한 처리가 완료되는 것을 attention 구현 단에서 sync하게 기다린다.

 

위에서 스크린샷을 다시 가져오자면, vLLM에서 pagedAttention의 forward는 아래와 같이 k, v, k_cache, v_cache를 인자로 전달 받는다.

 

이렇게 전달받은 k, v, k_cache, v_cache는 자체로 제작한 커널을 통해 reshape되고 concat된다.

 

아래는 reshape_and_cache에 대한 커널 구현체이다. 커널을 이해하기 위해서는 cuda 프로그래밍에 대한 이해가 있어야 하므로 이 포스트에서는 다루지 않겠다.

 

이렇게 구해진 kv_cache는 마찬가지로 자체 구현된 attention op 커널에 의해 output값을 채운다.

 

정리하자면, vLLM은 자체적으로 cache를 관리하는 엔진을 별도로 구현했으며 huggingface transformers와 달리 cpu와 gpu를 동시에 사용할 수 있는 메모리적 이점이 있다. 또한 cache와 관련된 여러 operation들을 하나의 커널로 자체 구현하여 속도를 향상시켰음을 알 수 있다.

 

---

 

이상으로 본 포스트에서는 KV Cache에 대해서 좀 더 deep dive하여, llama 모델의 구현체 위주로 huggingface transformers의 구현체와 vLLM 프로젝트의 구현체를 살펴보았다. vLLM 프로젝트의 경우에는 KV Cache 및 메모리에 좀 더 중점을 맞춰서 나온 새로운 프로젝트라서 그런지 huggingface transformers보다 많은 고민들을 한 모습이 돋보였다.

다음 포스트에서는 다시 본론으로 돌아가서 좀 더 다양한 측면에서의 가속 방법론을 다룰 생각이다.