Inference와 KV Cache 최적화 완벽 가이드

대규모 언어 모델의 추론 속도를 높이는 핵심 기술인 KV Cache의 원리부터 배치 추론 최적화까지 다룹니다. 초급 개발자도 이해할 수 있도록 실무 예제와 함께 설명합니다.

---

목차

1. KV_Cache가_필요한_이유
2. engine_py_소스_코드_분석
3. 토큰_생성_루프_구현
4. Temperature와_Top_p_샘플링
5. 배치_추론_최적화
6. 메모리_vs_속도_트레이드오프

---

1. KV Cache가 필요한 이유

어느 날 김개발 씨는 회사에서 LLM 기반 챗봇 서비스를 개발하고 있었습니다. 그런데 이상하게도 응답이 너무 느렸습니다.

한 문장을 생성하는 데 10초가 넘게 걸리는 것이었습니다. 선배 박시니어 씨가 다가와 물었습니다.

"혹시 KV Cache 안 쓰고 있어요?"

KV Cache는 한마디로 이전에 계산한 결과를 저장해두고 재사용하는 기술입니다. 마치 시험 볼 때 앞에서 풀었던 문제의 풀이 과정을 메모해두고, 비슷한 문제가 나오면 다시 계산하지 않고 참고하는 것과 같습니다.

이것을 제대로 활용하면 LLM의 추론 속도를 수십 배까지 높일 수 있습니다.

다음 코드를 살펴봅시다.

# KV Cache 없이 토큰 생성 (비효율적)
def generate_without_cache(model, input_ids, max_length):
    for i in range(max_length):
        # 매번 전체 시퀀스를 다시 계산합니다
        outputs = model(input_ids)  # O(n^2) 복잡도!
        next_token = outputs.logits[:, -1, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)
    return input_ids

# KV Cache를 활용한 토큰 생성 (효율적)
def generate_with_cache(model, input_ids, max_length):
    past_key_values = None  # KV Cache 저장소
    for i in range(max_length):
        # 새 토큰만 계산하고 이전 결과는 캐시에서 가져옵니다
        outputs = model(input_ids[:, -1:], past_key_values=past_key_values)
        past_key_values = outputs.past_key_values  # 캐시 업데이트
        next_token = outputs.logits[:, -1, :].argmax(dim=-1)
        input_ids = torch.cat([input_ids, next_token.unsqueeze(-1)], dim=-1)
    return input_ids

김개발 씨는 입사 6개월 차 주니어 개발자입니다. 요즘 회사에서 가장 핫한 프로젝트인 AI 챗봇 개발에 참여하게 되어 설레는 마음으로 코드를 작성하고 있었습니다.

그런데 문제가 생겼습니다. 사용자가 질문을 하면 답변이 나오기까지 한참을 기다려야 했던 것입니다.

"이상하다, 분명 GPU도 좋은 걸 쓰고 있는데..." 김개발 씨가 혼잣말을 하자, 옆자리의 박시니어 씨가 모니터를 슬쩍 보더니 웃음을 지었습니다. "아, KV Cache를 안 쓰고 있구나." 그렇다면 KV Cache란 정확히 무엇일까요?

쉽게 비유하자면, KV Cache는 마치 요리사가 미리 손질해둔 재료 보관함과 같습니다. 손님이 파스타를 주문할 때마다 매번 양파를 껍질부터 까고, 마늘을 하나하나 다지는 것은 너무 비효율적입니다.

똑똑한 요리사는 미리 재료를 손질해서 보관해두고, 주문이 들어오면 바로 꺼내 씁니다. KV Cache도 마찬가지로, 이미 계산한 Key와 Value 값을 저장해두고 재사용합니다.

KV Cache가 없던 시절에는 어땠을까요? Transformer 모델이 토큰을 하나 생성할 때마다, 이전에 생성한 모든 토큰에 대해 어텐션 연산을 다시 수행해야 했습니다.

예를 들어 "오늘 날씨가 좋습니다"라는 문장을 생성한다고 해봅시다. "오늘"을 생성할 때 1번 계산하고, "날씨가"를 생성할 때는 "오늘"부터 다시 2번 계산합니다.

"좋습니다"를 생성할 때는 또다시 처음부터 3번 계산해야 합니다. 토큰이 늘어날수록 계산량은 기하급수적으로 증가합니다.

이것이 바로 O(n^2) 복잡도의 함정입니다. 100개의 토큰을 생성하려면 1+2+3+...+100, 즉 5,050번의 계산이 필요합니다.

1,000개의 토큰이라면 50만 번이 넘습니다. 바로 이런 문제를 해결하기 위해 KV Cache가 등장했습니다.

KV Cache를 사용하면 이전 토큰들의 Key와 Value 값을 메모리에 저장해둡니다. 새 토큰을 생성할 때는 저장된 캐시를 그대로 가져다 쓰고, 새 토큰에 대한 계산만 추가로 수행합니다.

이렇게 하면 복잡도가 **O(n)**으로 줄어듭니다. 100개 토큰을 생성해도 100번만 계산하면 됩니다.

위의 코드를 자세히 살펴보겠습니다. 캐시 없는 버전에서는 매 반복마다 model(input_ids)를 호출합니다.

이때 input_ids의 길이가 계속 늘어나므로, 연산량도 함께 증가합니다. 반면 캐시를 사용하는 버전에서는 past_key_values 파라미터를 통해 이전 계산 결과를 전달합니다.

모델은 마지막 토큰(input_ids[:, -1:])에 대해서만 새로 계산하고, 나머지는 캐시에서 가져옵니다. 실제 현업에서는 어떻게 활용할까요?

Hugging Face의 transformers 라이브러리는 기본적으로 KV Cache를 지원합니다. model.generate() 메서드를 호출하면 내부적으로 캐시를 활용한 최적화가 자동으로 적용됩니다.

하지만 커스텀 추론 파이프라인을 구축할 때는 직접 캐시를 관리해야 하는 경우도 있습니다. 하지만 주의할 점도 있습니다.

KV Cache는 메모리를 상당히 많이 사용합니다. 긴 문맥을 처리하면 캐시 크기가 기가바이트 단위로 커질 수 있습니다.

따라서 메모리가 부족한 환경에서는 캐시 크기를 제한하거나, 오래된 캐시를 삭제하는 전략이 필요합니다. 다시 김개발 씨의 이야기로 돌아가 봅시다.

박시니어 씨의 설명을 들은 김개발 씨는 코드에 KV Cache를 적용했습니다. 결과는 놀라웠습니다.

10초 걸리던 응답이 1초도 안 되어 나왔습니다. "와, 이렇게 간단한 것으로 이렇게 빨라지다니!" 김개발 씨는 감탄했습니다.

실전 팁

💡 - KV Cache는 메모리와 속도의 트레이드오프입니다. GPU 메모리 상황을 항상 모니터링하세요.

• 배치 추론 시에는 각 샘플마다 별도의 캐시를 관리해야 합니다.
• 긴 문맥에서는 Sliding Window Attention과 함께 사용하면 효과적입니다.

---

2. engine py 소스 코드 분석

박시니어 씨는 김개발 씨에게 회사의 추론 엔진 코드를 보여주었습니다. "이게 우리 팀이 만든 engine.py야.

여기 보면 KV Cache가 어떻게 구현되어 있는지 알 수 있어." 김개발 씨는 눈을 반짝이며 코드를 살펴보기 시작했습니다.

engine.py는 LLM 추론의 핵심 로직을 담고 있는 파일입니다. 마치 자동차의 엔진처럼, 모델이 입력을 받아 출력을 생성하는 전체 과정을 관장합니다.

이 파일을 분석하면 KV Cache가 실제로 어떻게 초기화되고, 업데이트되며, 관리되는지 이해할 수 있습니다.

다음 코드를 살펴봅시다.

class InferenceEngine:
    def __init__(self, model, max_seq_len=2048):
        self.model = model
        self.max_seq_len = max_seq_len
        self.kv_cache = None  # KV Cache 저장소

    def initialize_cache(self, batch_size):
        # 모델 레이어 수와 헤드 수에 맞춰 캐시 초기화
        num_layers = self.model.config.num_hidden_layers
        num_heads = self.model.config.num_attention_heads
        head_dim = self.model.config.hidden_size // num_heads

        # (batch, num_heads, seq_len, head_dim) 형태로 초기화
        self.kv_cache = [
            (torch.zeros(batch_size, num_heads, 0, head_dim),
             torch.zeros(batch_size, num_heads, 0, head_dim))
            for _ in range(num_layers)
        ]

    def update_cache(self, new_keys, new_values, layer_idx):
        # 새로운 K, V를 기존 캐시에 연결
        old_k, old_v = self.kv_cache[layer_idx]
        self.kv_cache[layer_idx] = (
            torch.cat([old_k, new_keys], dim=2),
            torch.cat([old_v, new_values], dim=2)
        )

김개발 씨는 engine.py 파일을 열어보았습니다. 처음에는 복잡해 보였지만, 박시니어 씨의 설명을 들으며 하나씩 이해해 나갔습니다.

가장 먼저 눈에 들어온 것은 InferenceEngine 클래스였습니다. 이 클래스는 마치 공장의 생산 라인과 같습니다.

원재료(입력 토큰)가 들어오면, 여러 공정(모델 레이어)을 거쳐 완제품(출력 토큰)이 나옵니다. 그리고 각 공정에서 생산된 중간 부품(Key, Value)은 창고(KV Cache)에 보관됩니다.

__init__ 메서드를 살펴봅시다. 여기서는 모델과 최대 시퀀스 길이를 받아 저장합니다.

self.kv_cache = None으로 초기화하는데, 이는 아직 캐시가 비어있다는 의미입니다. 실제 추론이 시작될 때 initialize_cache 메서드가 호출되어 캐시가 생성됩니다.

initialize_cache 메서드가 캐시를 어떻게 만드는지 살펴봅시다. 먼저 모델의 설정에서 필요한 정보를 가져옵니다.

num_layers는 Transformer 레이어의 개수입니다. GPT-2 small은 12개, GPT-3는 96개의 레이어를 가지고 있습니다.

num_heads는 어텐션 헤드의 개수이고, head_dim은 각 헤드의 차원입니다. 캐시의 형태는 (batch, num_heads, seq_len, head_dim)입니다.

이 4차원 텐서가 의미하는 바를 하나씩 살펴보면, batch는 동시에 처리하는 샘플 수, num_heads는 병렬로 작동하는 어텐션 헤드 수, seq_len은 지금까지 처리한 토큰 수, head_dim은 각 헤드의 벡터 차원입니다. 처음에는 seq_len이 0으로 시작하고, 토큰이 생성될 때마다 증가합니다.

update_cache 메서드는 캐시를 업데이트합니다. 새로운 토큰에 대한 Key와 Value가 계산되면, 기존 캐시에 연결(concatenate)합니다.

torch.cat을 사용해 dim=2 방향, 즉 시퀀스 길이 방향으로 이어붙입니다. 이렇게 하면 캐시가 점점 길어지면서 모든 이전 토큰의 정보를 담게 됩니다.

실제 프로덕션 코드에서는 더 많은 기능이 필요합니다. 예를 들어 캐시가 최대 길이를 초과하면 오래된 토큰을 삭제하는 로직, 배치 내에서 각 샘플의 길이가 다를 때 처리하는 로직, GPU 메모리 부족 시 일부 캐시를 CPU로 옮기는 로직 등이 추가됩니다.

김개발 씨가 물었습니다. "그런데 왜 레이어마다 별도의 캐시가 필요한 거예요?" 박시니어 씨가 답했습니다.

"Transformer는 여러 레이어가 쌓여 있잖아. 각 레이어에서 어텐션을 계산할 때 그 레이어만의 Key, Value가 필요해.

그래서 레이어 수만큼 캐시를 만들어두는 거야." 김개발 씨는 고개를 끄덕였습니다. 코드의 의미가 하나씩 명확해지는 느낌이었습니다.

실전 팁

💡 - 캐시 초기화는 추론 시작 전에 한 번만 수행합니다. 매 토큰마다 초기화하면 안 됩니다.

• 멀티 GPU 환경에서는 캐시도 적절한 디바이스에 할당해야 합니다.
• 디버깅 시 캐시의 shape을 자주 확인하세요. 차원 불일치 오류가 흔합니다.

---

3. 토큰 생성 루프 구현

이론을 배운 김개발 씨는 이제 직접 코드를 작성해볼 차례였습니다. 박시니어 씨가 말했습니다.

"토큰 생성 루프를 직접 구현해봐. 그래야 진짜 이해할 수 있어." 김개발 씨는 키보드에 손을 올렸습니다.

토큰 생성 루프는 LLM이 텍스트를 한 토큰씩 순차적으로 생성하는 핵심 과정입니다. 마치 소설가가 한 단어씩 써내려가는 것처럼, 모델도 이전에 생성한 토큰들을 바탕으로 다음 토큰을 예측합니다.

이 루프를 효율적으로 구현하는 것이 추론 성능의 핵심입니다.

다음 코드를 살펴봅시다.

def generate_tokens(engine, input_ids, max_new_tokens=100, eos_token_id=None):
    batch_size = input_ids.shape[0]
    engine.initialize_cache(batch_size)

    # 첫 번째 forward pass: 프롬프트 전체 처리
    outputs = engine.model(input_ids, use_cache=True)
    engine.kv_cache = outputs.past_key_values
    next_token_logits = outputs.logits[:, -1, :]

    generated_tokens = []
    for step in range(max_new_tokens):
        # 다음 토큰 선택 (여기서는 greedy decoding)
        next_token = next_token_logits.argmax(dim=-1, keepdim=True)
        generated_tokens.append(next_token)

        # 종료 조건 확인
        if eos_token_id is not None and (next_token == eos_token_id).all():
            break

        # 캐시를 활용한 증분 추론
        outputs = engine.model(next_token, past_key_values=engine.kv_cache, use_cache=True)
        engine.kv_cache = outputs.past_key_values
        next_token_logits = outputs.logits[:, -1, :]

    return torch.cat(generated_tokens, dim=-1)

김개발 씨는 빈 파이썬 파일을 열고 generate_tokens 함수를 작성하기 시작했습니다. 먼저 전체 흐름을 생각해보았습니다.

토큰 생성은 크게 두 단계로 나뉩니다. 첫 번째는 프롬프트 처리 단계입니다.

사용자가 입력한 텍스트를 모델에 넣고 초기 KV Cache를 생성합니다. 두 번째는 토큰 생성 단계입니다.

루프를 돌면서 한 토큰씩 생성하고, 캐시를 업데이트합니다. input_ids.shape[0]으로 배치 크기를 확인합니다.

배치 크기를 알아야 캐시를 적절한 크기로 초기화할 수 있습니다. 만약 동시에 5개의 질문에 답변을 생성한다면, batch_size는 5가 됩니다.

첫 번째 forward pass가 가장 중요합니다. engine.model(input_ids, use_cache=True)를 호출하면, 모델은 프롬프트 전체를 한 번에 처리합니다.

이때 use_cache=True 옵션으로 KV Cache를 생성하도록 지시합니다. 반환된 outputs.past_key_values에는 모든 레이어의 Key, Value가 담겨 있습니다.

outputs.logits[:, -1, :]는 무엇을 의미할까요? logits는 모델이 예측한 다음 토큰의 확률 분포입니다.

[:, -1, :]로 마지막 토큰 위치의 출력만 가져옵니다. 왜냐하면 우리가 관심 있는 것은 프롬프트 다음에 올 토큰이기 때문입니다.

이제 생성 루프가 시작됩니다. for step in range(max_new_tokens) 루프 안에서, 매 반복마다 하나의 토큰을 생성합니다.

argmax를 사용해 가장 확률이 높은 토큰을 선택하는 것을 Greedy Decoding이라고 합니다. 가장 단순하지만 때로는 최적이 아닌 결과를 낼 수 있습니다.

종료 조건 확인도 중요합니다. EOS(End of Sequence) 토큰이 생성되면 더 이상 생성할 필요가 없습니다.

(next_token == eos_token_id).all()은 배치 내 모든 샘플이 종료되었는지 확인합니다. 증분 추론 부분을 자세히 봅시다.

engine.model(next_token, past_key_values=engine.kv_cache, use_cache=True)에서, 입력으로 새로 생성한 토큰 하나만 넣습니다. past_key_values로 이전 캐시를 전달하면, 모델은 이 캐시와 새 토큰을 결합해 다음 예측을 수행합니다.

이것이 KV Cache의 마법입니다. 마지막으로 생성된 토큰들을 하나로 합칩니다.

torch.cat(generated_tokens, dim=-1)로 모든 토큰을 연결해 최종 결과를 만듭니다. 이 결과를 디코딩하면 사람이 읽을 수 있는 텍스트가 됩니다.

김개발 씨는 코드를 완성하고 실행해보았습니다. 정상적으로 동작했습니다!

작은 성취감이 밀려왔습니다.

실전 팁

💡 - Greedy Decoding은 빠르지만 다양성이 부족합니다. 실제 서비스에서는 샘플링 기법을 사용하세요.

• max_new_tokens를 너무 크게 설정하면 메모리 부족이 발생할 수 있습니다.
• 배치 내 샘플들의 종료 시점이 다르면 패딩 처리가 필요합니다.

---

4. Temperature와 Top p 샘플링

김개발 씨가 만든 챗봇은 잘 동작했지만, 답변이 너무 딱딱했습니다. 같은 질문에 항상 똑같은 답만 내놓았습니다.

박시니어 씨가 말했습니다. "Greedy Decoding만 쓰면 그래.

Temperature랑 Top-p를 적용해봐."

Temperature와 Top-p는 LLM의 출력에 다양성을 부여하는 샘플링 기법입니다. Temperature는 마치 요리의 간 조절과 같아서, 낮으면 보수적이고 높으면 모험적인 답변이 나옵니다.

Top-p는 상위 몇 퍼센트의 후보만 고려할지 결정합니다. 이 두 기법을 잘 조합하면 창의적이면서도 일관성 있는 텍스트를 생성할 수 있습니다.

다음 코드를 살펴봅시다.

def sample_next_token(logits, temperature=1.0, top_p=0.9):
    # Temperature 적용: 확률 분포를 조절합니다
    if temperature != 1.0:
        logits = logits / temperature  # 높을수록 분포가 평평해짐

    # Softmax로 확률 변환
    probs = torch.softmax(logits, dim=-1)

    # Top-p (Nucleus) Sampling
    sorted_probs, sorted_indices = torch.sort(probs, descending=True)
    cumulative_probs = torch.cumsum(sorted_probs, dim=-1)

    # 누적 확률이 top_p를 넘는 토큰들을 제거
    sorted_indices_to_remove = cumulative_probs > top_p
    sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
    sorted_indices_to_remove[..., 0] = False  # 최소 1개는 유지

    # 제거할 토큰의 확률을 0으로 설정
    indices_to_remove = sorted_indices_to_remove.scatter(
        dim=-1, index=sorted_indices, src=sorted_indices_to_remove
    )
    probs[indices_to_remove] = 0.0
    probs = probs / probs.sum(dim=-1, keepdim=True)  # 재정규화

    # 확률에 따라 샘플링
    next_token = torch.multinomial(probs, num_samples=1)
    return next_token

김개발 씨는 Greedy Decoding의 한계를 체감하고 있었습니다. 사용자가 "재미있는 이야기 해줘"라고 하면, 챗봇은 항상 똑같은 이야기만 했습니다.

마치 녹음된 메시지를 재생하는 것 같았습니다. 사용자들은 금방 지루해했습니다.

Temperature가 어떻게 작동하는지 살펴봅시다. Temperature는 softmax 함수에 적용되는 스케일링 인자입니다.

logits를 temperature로 나누면, 확률 분포의 모양이 바뀝니다. temperature가 1.0보다 작으면 분포가 뾰족해져서 확률이 높은 토큰에 더욱 집중됩니다.

반대로 1.0보다 크면 분포가 평평해져서 다양한 토큰이 선택될 가능성이 높아집니다. 비유하자면, temperature는 룰렛 게임의 판 모양을 바꾸는 것과 같습니다.

낮은 temperature는 당첨 영역이 작은 룰렛입니다. 공이 거의 항상 같은 곳에 떨어집니다.

높은 temperature는 당첨 영역이 넓어진 룰렛입니다. 공이 어디에 떨어질지 예측하기 어렵습니다.

Top-p (또는 Nucleus Sampling)는 다른 접근법입니다. 확률이 높은 토큰들을 순서대로 나열하고, 누적 확률이 p를 넘을 때까지만 후보로 삼습니다.

예를 들어 top_p=0.9라면, 상위 90%의 확률 질량을 차지하는 토큰들만 고려합니다. 확률이 아주 낮은 이상한 토큰들은 아예 배제됩니다.

왜 둘 다 필요할까요? Temperature만 높이면 말도 안 되는 토큰이 선택될 수 있습니다.

"오늘 날씨가" 다음에 갑자기 "바나나"가 나올 수 있죠. Top-p는 이런 극단적인 선택을 막아줍니다.

두 기법을 함께 사용하면, 적당히 창의적이면서도 문맥에 맞는 텍스트가 생성됩니다. 코드를 단계별로 분석해봅시다.

먼저 logits / temperature로 스케일링합니다. 그 다음 softmax로 확률로 변환합니다.

torch.sort로 확률을 내림차순 정렬하고, torch.cumsum으로 누적 합을 구합니다. 누적 확률이 top_p를 넘는 토큰들은 마스킹해서 제거합니다.

마지막으로 torch.multinomial로 남은 토큰들 중에서 확률적으로 샘플링합니다. 실무에서 권장되는 설정값이 있습니다.

일반적인 대화에는 temperature=0.7, top_p=0.9 정도가 좋습니다. 코드 생성처럼 정확성이 중요한 작업에는 temperature=0.2로 낮춥니다.

창작 글쓰기에는 temperature=1.0 이상을 사용하기도 합니다. 김개발 씨는 샘플링을 적용한 후 챗봇을 테스트해보았습니다.

같은 질문에도 매번 조금씩 다른 답변이 나왔습니다. 훨씬 자연스러워 보였습니다.

실전 팁

💡 - temperature=0은 Greedy Decoding과 같습니다. 0에 가까울수록 결정적입니다.

• top_p=1.0은 모든 토큰을 후보로 삼는 것입니다. 사실상 비활성화와 같습니다.
• 서비스에 따라 적절한 값을 실험적으로 찾아야 합니다. 정답은 없습니다.

---

5. 배치 추론 최적화

챗봇 서비스가 인기를 얻으면서, 동시 접속자가 폭발적으로 늘어났습니다. 한 번에 한 명씩 처리하던 방식으로는 감당이 안 됐습니다.

박시니어 씨가 말했습니다. "이제 배치 추론을 도입해야 할 때야.

GPU를 제대로 활용하려면 필수야."

배치 추론은 여러 요청을 묶어서 한 번에 처리하는 기법입니다. 마치 버스가 승객을 한 명씩 태우지 않고 여러 명을 모아서 출발하는 것과 같습니다.

GPU는 병렬 연산에 특화되어 있어서, 배치 크기를 늘리면 전체 처리량(throughput)이 크게 향상됩니다.

다음 코드를 살펴봅시다.

class BatchInferenceEngine:
    def __init__(self, model, max_batch_size=32, max_seq_len=2048):
        self.model = model
        self.max_batch_size = max_batch_size
        self.request_queue = []
        self.kv_caches = {}  # request_id별 캐시 관리

    def add_request(self, request_id, input_ids):
        self.request_queue.append({
            'id': request_id,
            'input_ids': input_ids,
            'generated': [],
            'finished': False
        })

    def step(self):
        # 활성 요청만 필터링
        active_requests = [r for r in self.request_queue if not r['finished']]
        if not active_requests:
            return

        # 배치로 묶기 (패딩 처리 포함)
        batch_input_ids, attention_mask = self._prepare_batch(active_requests)
        batch_kv_cache = self._gather_caches(active_requests)

        # 배치 추론 수행
        outputs = self.model(
            batch_input_ids,
            attention_mask=attention_mask,
            past_key_values=batch_kv_cache,
            use_cache=True
        )

        # 결과 분배 및 캐시 업데이트
        self._distribute_outputs(active_requests, outputs)

김개발 씨의 챗봇 서비스는 대성공이었습니다. 하지만 성공에는 대가가 따랐습니다.

동시에 100명이 질문을 하면, 99명은 앞 사람이 끝날 때까지 기다려야 했습니다. 응답 시간이 몇 분씩 걸리는 경우도 있었습니다.

사용자들의 불만이 쏟아졌습니다. 배치 추론이 이 문제의 해결책입니다.

GPU는 행렬 연산에 최적화되어 있습니다. 하나의 벡터를 처리하나 32개의 벡터를 처리하나, 걸리는 시간이 거의 비슷합니다.

이 특성을 활용해 여러 요청을 하나의 배치로 묶어 처리하면, 전체 처리량이 크게 늘어납니다. 쉽게 비유하자면 세탁기와 같습니다.

옷 한 벌을 빨든 열 벌을 빨든, 세탁기 돌리는 시간은 비슷합니다. 옷을 한 벌씩 따로 빨면 하루 종일 걸리지만, 모아서 빨면 몇 시간이면 충분합니다.

GPU도 마찬가지입니다. 코드를 살펴봅시다.

BatchInferenceEngine 클래스는 여러 요청을 관리합니다. request_queue에 들어온 요청들을 저장하고, kv_caches에 각 요청별 캐시를 관리합니다.

add_request 메서드로 새 요청을 추가합니다. 각 요청에는 고유한 id, 입력 input_ids, 생성된 토큰들 generated, 완료 여부 finished가 포함됩니다.

요청이 들어올 때마다 큐에 추가됩니다. step 메서드가 핵심입니다.

매 호출마다 하나의 토큰을 생성합니다. 먼저 아직 완료되지 않은 활성 요청들만 필터링합니다.

그 다음 이들을 하나의 배치로 묶습니다. 요청마다 시퀀스 길이가 다를 수 있으므로 패딩 처리가 필요합니다.

_prepare_batch 메서드는 패딩을 처리합니다. 가장 긴 시퀀스에 맞춰 짧은 시퀀스들을 PAD 토큰으로 채웁니다.

attention_mask를 만들어 모델이 패딩 토큰을 무시하도록 합니다. 배치 추론을 수행한 후, 결과를 각 요청에 분배합니다.

_distribute_outputs 메서드에서 배치 출력을 개별 요청으로 나눕니다. 각 요청의 생성된 토큰 리스트에 새 토큰을 추가하고, 캐시도 업데이트합니다.

실제 프로덕션에서는 Continuous Batching을 사용합니다. 기존 방식은 배치 내 모든 요청이 끝날 때까지 기다려야 새 요청을 추가할 수 있었습니다.

Continuous Batching은 완료된 요청을 즉시 제거하고 새 요청을 추가합니다. 이렇게 하면 GPU 활용률이 더욱 높아집니다.

김개발 씨는 배치 추론을 적용한 후 동시 처리 능력이 10배 이상 향상된 것을 확인했습니다. 같은 GPU로 훨씬 많은 사용자를 감당할 수 있게 되었습니다.

실전 팁

💡 - 배치 크기가 너무 크면 메모리 부족이 발생합니다. GPU 메모리에 맞게 조절하세요.

• 요청마다 시퀀스 길이가 크게 다르면 패딩 오버헤드가 커집니다. 비슷한 길이끼리 묶는 것이 효율적입니다.
• vLLM, TensorRT-LLM 같은 프레임워크는 고급 배치 최적화를 기본 제공합니다.

---

6. 메모리 vs 속도 트레이드오프

모든 최적화가 적용된 시스템이 완성된 것 같았습니다. 하지만 박시니어 씨가 마지막 숙제를 내주었습니다.

"메모리와 속도 사이에는 항상 트레이드오프가 있어. 상황에 따라 어디에 무게를 둘지 결정해야 해."

메모리와 속도의 트레이드오프는 최적화의 영원한 딜레마입니다. KV Cache는 메모리를 사용해 속도를 얻는 대표적인 기법입니다.

반대로 메모리가 부족할 때는 캐시를 줄이거나 재계산해야 합니다. 서비스 요구사항에 따라 이 균형점을 찾는 것이 엔지니어의 역할입니다.

다음 코드를 살펴봅시다.

class AdaptiveKVCache:
    def __init__(self, model, memory_limit_gb=8):
        self.model = model
        self.memory_limit = memory_limit_gb * 1024 ** 3  # bytes로 변환
        self.cache_strategy = 'full'  # full, sliding, recompute

    def estimate_cache_size(self, batch_size, seq_len):
        # KV Cache 메모리 추정: 2 * layers * 2 * batch * heads * seq * head_dim * dtype_size
        config = self.model.config
        cache_size = (2 * config.num_hidden_layers * batch_size *
                      config.num_attention_heads * seq_len *
                      (config.hidden_size // config.num_attention_heads) * 2)  # float16
        return cache_size

    def get_cache_strategy(self, batch_size, max_seq_len):
        estimated_size = self.estimate_cache_size(batch_size, max_seq_len)

        if estimated_size < self.memory_limit * 0.5:
            return 'full'  # 전체 캐시 유지
        elif estimated_size < self.memory_limit * 0.8:
            return 'sliding'  # 슬라이딩 윈도우 캐시
        else:
            return 'recompute'  # 일부 레이어 재계산

    def apply_sliding_window(self, kv_cache, window_size=1024):
        # 최근 window_size 토큰만 캐시에 유지
        return [(k[:, :, -window_size:, :], v[:, :, -window_size:, :])
                for k, v in kv_cache]

김개발 씨는 서비스 운영 중에 가끔 이상한 현상을 발견했습니다. 특정 시간대에 GPU 메모리 사용량이 급격히 늘어나면서 서비스가 불안정해졌습니다.

긴 대화를 나누는 사용자가 많아지면, KV Cache가 메모리를 잡아먹었던 것입니다. 메모리 vs 속도 트레이드오프를 이해해야 이 문제를 해결할 수 있습니다.

비유하자면, 이것은 책상 위 공간과 같습니다. 책상이 넓으면 필요한 자료를 다 펼쳐놓고 작업할 수 있습니다.

하지만 책상이 좁으면 자료를 번갈아가며 꺼내야 합니다. 속도는 느리지만 작은 책상으로도 작업이 가능합니다.

KV Cache의 메모리 사용량을 계산해봅시다. estimate_cache_size 함수를 보면, 캐시 크기는 레이어 수, 배치 크기, 시퀀스 길이, 헤드 수, 헤드 차원에 비례합니다.

예를 들어 LLaMA-7B 모델(32 레이어, 32 헤드)로 배치 크기 8, 시퀀스 길이 4096을 처리하면, 약 16GB의 캐시가 필요합니다. 세 가지 캐시 전략이 있습니다.

full 전략은 모든 토큰의 캐시를 유지합니다. 가장 빠르지만 메모리를 가장 많이 씁니다.

sliding 전략은 최근 N개 토큰의 캐시만 유지합니다. 오래된 토큰 정보는 잃지만 메모리를 절약합니다.

recompute 전략은 일부 레이어의 캐시를 저장하지 않고 필요할 때 다시 계산합니다. get_cache_strategy 함수는 자동으로 전략을 선택합니다.

예상 캐시 크기가 메모리 한도의 50% 미만이면 full, 80% 미만이면 sliding, 그 이상이면 recompute를 선택합니다. 이렇게 적응적으로 전략을 바꾸면 메모리 부족 오류를 예방할 수 있습니다.

슬라이딩 윈도우 캐시를 자세히 살펴봅시다. apply_sliding_window 함수는 캐시에서 최근 window_size 토큰만 남기고 나머지를 버립니다.

k[:, :, -window_size:, :]는 마지막 window_size개 위치만 슬라이싱합니다. 이렇게 하면 캐시 크기가 window_size로 고정됩니다.

하지만 정보 손실이 있습니다. 아주 긴 대화에서 초반에 언급된 내용을 참조해야 할 때, 슬라이딩 윈도우로는 해당 정보에 접근할 수 없습니다.

따라서 서비스 특성에 맞게 window_size를 신중히 결정해야 합니다. 실무에서는 더 정교한 기법들이 사용됩니다.

PagedAttention은 vLLM에서 도입한 기법으로, 캐시를 페이지 단위로 관리해 메모리 단편화를 줄입니다. FlashAttention은 캐시를 청크로 나눠 처리해 메모리 접근 패턴을 최적화합니다.

Grouped Query Attention은 KV 헤드 수를 줄여 캐시 크기 자체를 감소시킵니다. 김개발 씨는 AdaptiveKVCache를 도입한 후, 메모리 사용량이 안정화되는 것을 확인했습니다.

트래픽 폭증 시에도 서비스가 안정적으로 동작했습니다. 박시니어 씨가 말했습니다.

"축하해, 이제 LLM 추론의 핵심을 다 배운 거야. 물론 배울 게 더 많지만, 기초는 탄탄해졌어." 김개발 씨는 뿌듯한 마음으로 고개를 끄덕였습니다.

KV Cache부터 시작해 배치 추론, 메모리 최적화까지. 길었지만 보람찬 여정이었습니다.

실전 팁

💡 - GPU 메모리 모니터링 도구(nvidia-smi, PyTorch Profiler)를 활용해 실시간 사용량을 파악하세요.

• 프로덕션에서는 메모리 여유분을 항상 확보해야 합니다. 100% 활용은 위험합니다.
• vLLM, TGI 같은 서빙 프레임워크는 이러한 트레이드오프를 자동으로 관리해줍니다.

---

이상으로 학습을 마칩니다. 위 내용을 직접 코드로 작성해보면서 익혀보세요!