LLM 추론은 계산 집약적이고 메모리를 많이 사용한다. 프로덕션 환경에서 LLM을 효율적으로 서빙하려면 다양한 최적화 기법이 필요하다. 이 글은 KV Cache, Quantization, Speculative Decoding 등 핵심 최적화 기술을 실전 예제와 함께 정리한다.
\nimport torch\nfrom transformers import AutoModelForCausalLM, AutoTokenizer\n\n# 모델 로드\nmodel = AutoModelForCausalLM.from_pretrained(\"gpt2\")\ntokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n\n# 기본 추론\ndef basic_inference(prompt: str, max_length: int = 100):\n inputs = tokenizer(prompt, return_tensors=\"pt\")\n \n # 각 토큰 생성마다 전체 모델을 통과\n outputs = model.generate(\n inputs.input_ids,\n max_length=max_length,\n do_sample=True,\n temperature=0.7\n )\n \n return tokenizer.decode(outputs[0], skip_special_tokens=True)LLM 추론의 주요 병목 지점:
\nKV Cache는 이전 토큰들의 Key와 Value를 캐싱하여 반복 계산을 피하는 기법이다.
\nimport torch\nimport torch.nn as nn\nfrom transformers import GPT2LMHeadModel\n\nclass OptimizedGPT2Inference:\n def __init__(self, model_name: str = \"gpt2\"):\n self.model = GPT2LMHeadModel.from_pretrained(model_name)\n self.tokenizer = AutoTokenizer.from_pretrained(model_name)\n self.kv_cache = None\n \n def generate_with_kv_cache(self, prompt: str, max_new_tokens: int = 50):\n inputs = self.tokenizer(prompt, return_tensors=\"pt\")\n input_ids = inputs.input_ids\n \n # 초기 KV Cache 생성\n past_key_values = None\n generated_ids = input_ids.clone()\n \n for _ in range(max_new_tokens):\n # KV Cache를 활용한 추론\n outputs = self.model(\n input_ids=generated_ids[:, -1:], # 마지막 토큰만 입력\n past_key_values=past_key_values, # 이전 KV Cache 사용\n use_cache=True\n )\n \n # 새로운 KV Cache 업데이트\n past_key_values = outputs.past_key_values\n \n # 다음 토큰 선택\n next_token_logits = outputs.logits[0, -1, :]\n next_token_id = torch.argmax(next_token_logits, dim=-1).unsqueeze(0)\n \n # 생성된 토큰 추가\n generated_ids = torch.cat([generated_ids, next_token_id.unsqueeze(0)], dim=1)\n \n return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)class KVCache:\n def __init__(self, num_layers: int, batch_size: int, seq_len: int, hidden_size: int, num_heads: int):\n self.num_layers = num_layers\n self.cache = {}\n \n # 각 레이어별 Key, Value 캐시 초기화\n for layer_idx in range(num_layers):\n self.cache[f\"layer_{layer_idx}_key\"] = torch.zeros(\n batch_size, num_heads, seq_len, hidden_size // num_heads\n )\n self.cache[f\"layer_{layer_idx}_value\"] = torch.zeros(\n batch_size, num_heads, seq_len, hidden_size // num_heads\n )\n \n def update(self, layer_idx: int, new_key: torch.Tensor, new_value: torch.Tensor):\n \"\"\"새로운 Key, Value로 캐시 업데이트\"\"\"\n old_key = self.cache[f\"layer_{layer_idx}_key\"]\n old_value = self.cache[f\"layer_{layer_idx}_value\"]\n \n # 기존 캐시와 새로운 값 연결\n self.cache[f\"layer_{layer_idx}_key\"] = torch.cat([old_key, new_key], dim=2)\n self.cache[f\"layer_{layer_idx}_value\"] = torch.cat([old_value, new_value], dim=2)\n \n def get(self, layer_idx: int):\n \"\"\"특정 레이어의 Key, Value 반환\"\"\"\n return (\n self.cache[f\"layer_{layer_idx}_key\"],\n self.cache[f\"layer_{layer_idx}_value\"]\n )import time\n\ndef benchmark_inference(model, tokenizer, prompt: str, use_cache: bool = True):\n inputs = tokenizer(prompt, return_tensors=\"pt\")\n \n start_time = time.time()\n \n if use_cache:\n # KV Cache 사용\n past_key_values = None\n generated_ids = inputs.input_ids\n \n for _ in range(50):\n outputs = model(\n input_ids=generated_ids[:, -1:],\n past_key_values=past_key_values,\n use_cache=True\n )\n past_key_values = outputs.past_key_values\n next_token = torch.argmax(outputs.logits[0, -1, :]).unsqueeze(0)\n generated_ids = torch.cat([generated_ids, next_token.unsqueeze(0)], dim=1)\n else:\n # KV Cache 미사용 (전체 시퀀스 재계산)\n generated_ids = inputs.input_ids\n for _ in range(50):\n outputs = model(input_ids=generated_ids)\n next_token = torch.argmax(outputs.logits[0, -1, :]).unsqueeze(0)\n generated_ids = torch.cat([generated_ids, next_token.unsqueeze(0)], dim=1)\n \n elapsed_time = time.time() - start_time\n return elapsed_time, generated_ids\n\n# 성능 비교\nprompt = \"The future of AI is\"\ntime_with_cache, _ = benchmark_inference(model, tokenizer, prompt, use_cache=True)\ntime_without_cache, _ = benchmark_inference(model, tokenizer, prompt, use_cache=False)\n\nprint(f\"KV Cache 사용: {time_with_cache:.2f}초\")\nprint(f\"KV Cache 미사용: {time_without_cache:.2f}초\")\nprint(f\"속도 향상: {time_without_cache / time_with_cache:.2f}x\")양자화는 모델의 가중치와 활성화 값을 낮은 비트로 변환하여 메모리 사용량과 계산 속도를 개선하는 기법이다.
\nfrom transformers import AutoModelForCausalLM, BitsAndBytesConfig\nimport torch\n\n# 8-bit 양자화 설정\nquantization_config = BitsAndBytesConfig(\n load_in_8bit=True,\n llm_int8_threshold=6.0,\n llm_int8_has_fp16_weight=False\n)\n\n# 양자화된 모델 로드\nmodel_8bit = AutoModelForCausalLM.from_pretrained(\n \"gpt2\",\n quantization_config=quantization_config,\n device_map=\"auto\"\n)\n\n# 4-bit 양자화 설정\nquantization_config_4bit = BitsAndBytesConfig(\n load_in_4bit=True,\n bnb_4bit_compute_dtype=torch.float16,\n bnb_4bit_use_double_quant=True,\n bnb_4bit_quant_type=\"nf4\"\n)\n\nmodel_4bit = AutoModelForCausalLM.from_pretrained(\n \"gpt2\",\n quantization_config=quantization_config_4bit,\n device_map=\"auto\"\n)import torch.quantization as quant\n\n# 모델 준비\nmodel.eval()\n\n# 양자화 설정\nmodel.qconfig = quant.get_default_qconfig('fbgemm')\n\n# 양자화 준비\nquant.prepare(model, inplace=True)\n\n# 캘리브레이션 데이터로 양자화\ncalibration_data = torch.randn(100, 10)\nwith torch.no_grad():\n for i in range(10):\n model(calibration_data)\n\n# 양자화 변환\nquant.convert(model, inplace=True)from auto_gptq import AutoGPTQForCausalLM, BaseQuantizeConfig\n\n# GPTQ 양자화 설정\nquantize_config = BaseQuantizeConfig(\n bits=4,\n group_size=128,\n desc_act=False\n)\n\n# 모델 양자화\nmodel = AutoGPTQForCausalLM.from_pretrained(\n \"gpt2\",\n quantize_config=quantize_config\n)\n\n# 양자화 실행\nexamples = [\n tokenizer(\"Hello, how are you?\", return_tensors=\"pt\")\n for _ in range(100)\n]\n\nmodel.quantize(examples)\n\n# 양자화된 모델 저장\nmodel.save_quantized(\"./gpt2-4bit\")import torch\n\ndef measure_memory_usage(model, input_ids):\n torch.cuda.reset_peak_memory_stats()\n _ = model(input_ids)\n memory_used = torch.cuda.max_memory_allocated() / 1024**3 # GB\n return memory_used\n\n# 원본 모델\noriginal_model = AutoModelForCausalLM.from_pretrained(\"gpt2\").cuda()\noriginal_memory = measure_memory_usage(original_model, input_ids)\n\n# 8-bit 양자화 모델\nmodel_8bit = AutoModelForCausalLM.from_pretrained(\n \"gpt2\",\n quantization_config=quantization_config,\n device_map=\"auto\"\n)\nmemory_8bit = measure_memory_usage(model_8bit, input_ids)\n\n# 4-bit 양자화 모델\nmodel_4bit = AutoModelForCausalLM.from_pretrained(\n \"gpt2\",\n quantization_config=quantization_config_4bit,\n device_map=\"auto\"\n)\nmemory_4bit = measure_memory_usage(model_4bit, input_ids)\n\nprint(f\"원본 모델: {original_memory:.2f} GB\")\nprint(f\"8-bit 양자화: {memory_8bit:.2f} GB ({original_memory/memory_8bit:.2f}x 감소)\")\nprint(f\"4-bit 양자화: {memory_4bit:.2f} GB ({original_memory/memory_4bit:.2f}x 감소)\")Speculative Decoding은 작은 모델로 여러 토큰을 예측하고, 큰 모델로 검증하여 병렬 처리를 가능하게 하는 기법이다.
\nclass SpeculativeDecoder:\n def __init__(self, draft_model, target_model, tokenizer):\n self.draft_model = draft_model # 작은 모델\n self.target_model = target_model # 큰 모델\n self.tokenizer = tokenizer\n \n def generate(self, prompt: str, max_tokens: int = 50, gamma: int = 4):\n \"\"\"\n gamma: 추측할 토큰 수\n \"\"\"\n input_ids = self.tokenizer(prompt, return_tensors=\"pt\").input_ids\n generated_ids = input_ids.clone()\n \n while len(generated_ids[0]) < max_tokens:\n # 1. Draft 모델로 gamma개의 토큰 추측\n draft_outputs = self.draft_model.generate(\n input_ids=generated_ids,\n max_new_tokens=gamma,\n do_sample=False,\n use_cache=True\n )\n draft_tokens = draft_outputs[0][len(generated_ids[0]):]\n \n # 2. Target 모델로 검증\n accepted_tokens = []\n for i, token in enumerate(draft_tokens):\n # Target 모델로 확률 계산\n target_outputs = self.target_model(\n input_ids=torch.cat([generated_ids, torch.tensor([accepted_tokens + [token]])], dim=1),\n use_cache=True\n )\n \n # Accept/Reject 결정\n draft_prob = self._get_draft_probability(draft_tokens[:i+1], generated_ids)\n target_prob = torch.softmax(target_outputs.logits[0, -1, :], dim=-1)[token]\n \n if self._should_accept(draft_prob, target_prob):\n accepted_tokens.append(token)\n else:\n # Reject 시 Target 모델에서 직접 샘플링\n new_token = torch.multinomial(target_prob, 1)\n accepted_tokens.append(new_token.item())\n break\n \n # 3. 수락된 토큰 추가\n generated_ids = torch.cat([\n generated_ids,\n torch.tensor([accepted_tokens]).unsqueeze(0)\n ], dim=1)\n \n return self.tokenizer.decode(generated_ids[0], skip_special_tokens=True)\n \n def _get_draft_probability(self, tokens, prefix):\n # Draft 모델의 확률 계산\n pass\n \n def _should_accept(self, draft_prob, target_prob):\n # Accept/Reject 결정 로직\n return target_prob >= draft_probimport torch.nn.functional as F\n\ndef parallel_speculative_decode(\n draft_model,\n target_model,\n input_ids,\n gamma: int = 4,\n temperature: float = 1.0\n):\n \"\"\"\n 병렬로 여러 토큰을 추측하고 검증\n \"\"\"\n # Draft 모델로 gamma개 토큰 생성\n with torch.no_grad():\n draft_outputs = draft_model.generate(\n input_ids,\n max_new_tokens=gamma,\n do_sample=True,\n temperature=temperature,\n use_cache=True\n )\n \n draft_sequence = draft_outputs[0][len(input_ids[0]):]\n \n # Target 모델로 전체 시퀀스 검증\n full_sequence = torch.cat([input_ids, draft_sequence.unsqueeze(0)], dim=1)\n target_outputs = target_model(full_sequence, use_cache=True)\n \n # 각 위치에서 Accept/Reject 결정\n accepted_tokens = []\n for i in range(len(draft_sequence)):\n draft_token = draft_sequence[i]\n \n # Target 모델의 확률\n target_logits = target_outputs.logits[0, len(input_ids[0]) + i - 1, :]\n target_probs = F.softmax(target_logits / temperature, dim=-1)\n \n # Draft 모델의 확률\n draft_logits = draft_outputs.logits[0, len(input_ids[0]) + i - 1, :]\n draft_probs = F.softmax(draft_logits / temperature, dim=-1)\n \n # Accept 확률 계산\n accept_prob = min(1.0, target_probs[draft_token] / draft_probs[draft_token])\n \n if torch.rand(1) < accept_prob:\n accepted_tokens.append(draft_token.item())\n else:\n # Reject: Target 모델에서 재샘플링\n new_token = torch.multinomial(target_probs, 1)\n accepted_tokens.append(new_token.item())\n break\n \n return torch.tensor([accepted_tokens])class DynamicBatching:\n def __init__(self, model, tokenizer, max_batch_size: int = 8):\n self.model = model\n self.tokenizer = tokenizer\n self.max_batch_size = max_batch_size\n self.request_queue = []\n \n def add_request(self, prompt: str, max_tokens: int = 50):\n \"\"\"요청을 큐에 추가\"\"\"\n self.request_queue.append({\n \"prompt\": prompt,\n \"max_tokens\": max_tokens,\n \"input_ids\": self.tokenizer(prompt, return_tensors=\"pt\").input_ids\n })\n \n def process_batch(self):\n \"\"\"배치 처리\"\"\"\n if not self.request_queue:\n return []\n \n # 배치 크기 결정\n batch_size = min(len(self.request_queue), self.max_batch_size)\n batch = self.request_queue[:batch_size]\n self.request_queue = self.request_queue[batch_size:]\n \n # 패딩으로 배치 구성\n max_length = max(len(req[\"input_ids\"][0]) for req in batch)\n batch_input_ids = []\n attention_masks = []\n \n for req in batch:\n input_ids = req[\"input_ids\"][0]\n padding_length = max_length - len(input_ids)\n \n padded_input = torch.cat([\n input_ids,\n torch.zeros(padding_length, dtype=torch.long)\n ])\n batch_input_ids.append(padded_input)\n \n attention_mask = torch.cat([\n torch.ones(len(input_ids)),\n torch.zeros(padding_length)\n ])\n attention_masks.append(attention_mask)\n \n batch_input_ids = torch.stack(batch_input_ids)\n attention_masks = torch.stack(attention_masks)\n \n # 배치 추론\n with torch.no_grad():\n outputs = self.model.generate(\n batch_input_ids,\n attention_mask=attention_masks,\n max_new_tokens=50,\n use_cache=True,\n pad_token_id=self.tokenizer.pad_token_id\n )\n \n # 결과 디코딩\n results = []\n for i, output in enumerate(outputs):\n decoded = self.tokenizer.decode(output, skip_special_tokens=True)\n results.append(decoded)\n \n return resultsclass ContinuousBatching:\n def __init__(self, model, tokenizer):\n self.model = model\n self.tokenizer = tokenizer\n self.active_requests = []\n self.kv_caches = {}\n \n def add_request(self, request_id: str, prompt: str, max_tokens: int):\n \"\"\"새 요청 추가\"\"\"\n input_ids = self.tokenizer(prompt, return_tensors=\"pt\").input_ids\n self.active_requests.append({\n \"id\": request_id,\n \"input_ids\": input_ids,\n \"max_tokens\": max_tokens,\n \"generated_tokens\": 0,\n \"kv_cache\": None\n })\n \n def process_step(self):\n \"\"\"한 스텝 처리\"\"\"\n if not self.active_requests:\n return\n \n # 모든 요청의 다음 토큰 생성\n for req in self.active_requests:\n if req[\"generated_tokens\"] >= req[\"max_tokens\"]:\n continue\n \n # 마지막 토큰만 사용\n last_token = req[\"input_ids\"][:, -1:]\n \n outputs = self.model(\n input_ids=last_token,\n past_key_values=req[\"kv_cache\"],\n use_cache=True\n )\n \n # KV Cache 업데이트\n req[\"kv_cache\"] = outputs.past_key_values\n \n # 다음 토큰 선택\n next_token = torch.argmax(outputs.logits[0, -1, :], dim=-1)\n req[\"input_ids\"] = torch.cat([req[\"input_ids\"], next_token.unsqueeze(0)], dim=1)\n req[\"generated_tokens\"] += 1\n \n # 완료된 요청 제거\n self.active_requests = [\n req for req in self.active_requests\n if req[\"generated_tokens\"] < req[\"max_tokens\"]\n ]import torch.nn as nn\nfrom torch.distributed import Pipeline\n\nclass PipelineParallelModel(nn.Module):\n def __init__(self, num_layers: int, hidden_size: int, num_gpus: int = 4):\n super().__init__()\n self.num_gpus = num_gpus\n self.layers_per_gpu = num_layers // num_gpus\n \n # 각 GPU에 할당될 레이어\n self.gpu_layers = nn.ModuleList([\n nn.ModuleList([\n nn.TransformerEncoderLayer(hidden_size, 8)\n for _ in range(self.layers_per_gpu)\n ])\n for _ in range(num_gpus)\n ])\n \n def forward(self, x):\n # Pipeline으로 순차 처리\n for gpu_idx in range(self.num_gpus):\n layers = self.gpu_layers[gpu_idx]\n for layer in layers:\n x = layer(x)\n return xclass TensorParallelAttention(nn.Module):\n def __init__(self, hidden_size: int, num_heads: int, num_gpus: int = 2):\n super().__init__()\n self.num_gpus = num_gpus\n self.num_heads_per_gpu = num_heads // num_gpus\n \n # 각 GPU에 할당될 Attention 헤드\n self.qkv_layers = nn.ModuleList([\n nn.Linear(hidden_size, hidden_size // num_gpus)\n for _ in range(num_gpus)\n ])\n \n def forward(self, x):\n # 각 GPU에서 부분 계산\n qkv_outputs = []\n for layer in self.qkv_layers:\n qkv_outputs.append(layer(x))\n \n # 결과 병합\n return torch.cat(qkv_outputs, dim=-1)try:\n from flash_attn import flash_attn_func\nexcept ImportError:\n print(\"Flash Attention이 설치되지 않았습니다.\")\n\ndef flash_attention_forward(q, k, v, dropout_p=0.0, softmax_scale=None):\n \"\"\"\n Flash Attention을 사용한 효율적인 Attention 계산\n \"\"\"\n return flash_attn_func(\n q, k, v,\n dropout_p=dropout_p,\n softmax_scale=softmax_scale,\n causal=True\n )\n\n# 사용 예시\nclass FlashAttentionLayer(nn.Module):\n def __init__(self, hidden_size: int, num_heads: int):\n super().__init__()\n self.hidden_size = hidden_size\n self.num_heads = num_heads\n self.head_dim = hidden_size // num_heads\n \n self.q_proj = nn.Linear(hidden_size, hidden_size)\n self.k_proj = nn.Linear(hidden_size, hidden_size)\n self.v_proj = nn.Linear(hidden_size, hidden_size)\n self.o_proj = nn.Linear(hidden_size, hidden_size)\n \n def forward(self, x):\n batch_size, seq_len, _ = x.shape\n \n q = self.q_proj(x)\n k = self.k_proj(x)\n v = self.v_proj(x)\n \n # Flash Attention 사용\n q = q.view(batch_size, seq_len, self.num_heads, self.head_dim)\n k = k.view(batch_size, seq_len, self.num_heads, self.head_dim)\n v = v.view(batch_size, seq_len, self.num_heads, self.head_dim)\n \n output = flash_attention_forward(q, k, v)\n output = output.contiguous().view(batch_size, seq_len, self.hidden_size)\n \n return self.o_proj(output)from fastapi import FastAPI, Request\nfrom pydantic import BaseModel\nimport torch\n\napp = FastAPI()\n\nclass InferenceRequest(BaseModel):\n prompt: str\n max_tokens: int = 50\n temperature: float = 0.7\n\nclass OptimizedLLMService:\n def __init__(self):\n # 양자화된 모델 로드\n self.model = AutoModelForCausalLM.from_pretrained(\n \"gpt2\",\n quantization_config=quantization_config_4bit,\n device_map=\"auto\"\n )\n self.tokenizer = AutoTokenizer.from_pretrained(\"gpt2\")\n self.batcher = ContinuousBatching(self.model, self.tokenizer)\n \n def generate(self, request: InferenceRequest):\n # 배치 처리로 효율적 생성\n request_id = str(uuid.uuid4())\n self.batcher.add_request(\n request_id,\n request.prompt,\n request.max_tokens\n )\n \n # 배치 처리\n result = self.batcher.process_step()\n return result\n\nservice = OptimizedLLMService()\n\n@app.post(\"/generate\")\nasync def generate(request: InferenceRequest):\n result = service.generate(request)\n return {\"generated_text\": result}def benchmark_all_techniques(model_name: str = \"gpt2\"):\n results = {}\n \n # 1. 기본 추론\n baseline_time = benchmark_baseline(model_name)\n results[\"baseline\"] = baseline_time\n \n # 2. KV Cache\n kv_cache_time = benchmark_kv_cache(model_name)\n results[\"kv_cache\"] = kv_cache_time\n results[\"kv_cache_speedup\"] = baseline_time / kv_cache_time\n \n # 3. 8-bit 양자화\n quant8_time, quant8_memory = benchmark_quantization(model_name, bits=8)\n results[\"quant8\"] = {\"time\": quant8_time, \"memory\": quant8_memory}\n \n # 4. 4-bit 양자화\n quant4_time, quant4_memory = benchmark_quantization(model_name, bits=4)\n results[\"quant4\"] = {\"time\": quant4_time, \"memory\": quant4_memory}\n \n # 5. Speculative Decoding\n spec_time = benchmark_speculative(model_name)\n results[\"speculative\"] = spec_time\n results[\"speculative_speedup\"] = baseline_time / spec_time\n \n return resultsLLM 추론 최적화는 다양한 기법을 조합하여 수행한다. 이 글에서 다룬 내용:
\n이러한 기법들을 적절히 조합하면 프로덕션 환경에서 LLM을 효율적으로 서빙할 수 있다.
\nLLM 추론 최적화 - KV Cache, Quantization, Speculative Decoding
\n