beautifulsoups.dev
Published on

로컬 LLM 느릴 때 - 4bit+FlashAttention 가속

Authors

로컬에서 transformers로 LLM을 띄웠는데 VRAM은 간당간당하고, 토큰 생성 속도는 1~5 tok/s 수준에서 답답하게 느려지는 경우가 많습니다. 이때 흔히 “양자화(4bit)하면 빨라지겠지”라고 생각하지만, 4bit는 주로 메모리 사용량을 줄이는 옵션이고 속도는 환경에 따라 오히려 비슷하거나 느려질 수도 있습니다.

그럼에도 불구하고 로컬 LLM을 확실히 체감 가속시키는 조합이 있습니다.

  • bitsandbytes 4bit로 VRAM을 확보해서 OOM을 피하고
  • FlashAttention 또는 PyTorch SDPA로 어텐션 연산을 빠른 커널로 교체하고
  • KV 캐시, prefill/decode 구간을 분리해 병목을 제거하는 방식입니다.

이 글은 “느린 원인을 분해”하고, 가장 흔한 실전 설정을 코드로 제시합니다. OOM 자체가 문제라면 먼저 Transformers 로컬 LLM OOM 해결 - 4bit+Offload도 함께 보시면 좋습니다.

왜 느린가: 로컬 LLM 성능을 결정하는 4가지 병목

1) 어텐션 커널이 느리다 (가장 흔함)

기본 어텐션 구현은 메모리 접근이 비효율적이라, 같은 GPU에서도 커널에 따라 속도가 크게 갈립니다.

  • PyTorch 2.x: scaled_dot_product_attention가 조건을 만족하면 내부적으로 빠른 커널을 선택합니다.
  • FlashAttention-2: 가능한 환경에서 더 공격적으로 최적화된 커널을 사용합니다.

즉, 모델이 느린 게 아니라 어텐션 경로가 느린 경우가 많습니다.

2) VRAM이 부족해 offload/스와핑이 발생한다

VRAM이 부족하면 다음이 발생합니다.

  • 일부 레이어가 CPU로 offload
  • paged 메모리, PCIe 왕복, CPU RAM 병목

이때 토큰 생성이 “뚝뚝 끊기는” 느낌이 납니다. 4bit는 이 문제를 완화하는 데 매우 효과적입니다.

3) prefill(프롬프트 처리)과 decode(토큰 생성)의 성격이 다르다

  • prefill: 입력 길이에 비례해 한 번에 큰 행렬 연산이 발생
  • decode: 토큰을 한 개씩 생성하며 KV 캐시를 계속 참조

프롬프트가 길수록 첫 응답까지 시간이 길어지고, 생성 속도(tok/s)만 보고 판단하면 원인을 놓치기 쉽습니다.

4) 설정 실수: 캐시 비활성화, dtype 불일치, CPU 토크나이저 병목

  • use_cache=False면 매 토큰마다 과거 토큰을 다시 계산해 급격히 느려집니다.
  • dtype가 꼬여서 불필요한 캐스팅이 발생하면 성능이 떨어집니다.
  • 토크나이저가 단일 스레드로 CPU에서 병목이 생기기도 합니다.

목표 전략: 4bit로 “들어가게” 만들고, FlashAttention로 “빠르게” 만든다

정리하면 접근은 다음 순서가 안전합니다.

  1. 4bit로 VRAM 확보 (OOM 방지 + offload 최소화)
  2. 어텐션 커널 가속 (FlashAttention 또는 SDPA)
  3. KV 캐시/생성 설정 정리 (불필요한 재계산 제거)
  4. 측정: prefill latency와 decode tok/s를 분리해서 본다

실전 코드: bitsandbytes 4bit + FlashAttention(또는 SDPA)

아래 예시는 가장 많이 쓰는 형태입니다.

  • bitsandbytes 4bit: nf4 + double quant
  • compute dtype: 보통 bfloat16 권장(지원 GPU일 때)
  • attention: 모델/버전에 따라 flash_attention_2 또는 sdpa

주의: 이 글은 MDX 렌더링 환경을 가정하므로, 본문에서 부등호 문자를 직접 노출하지 않습니다. 코드 블록 안에서는 안전합니다.

import torch
from transformers import AutoTokenizer, AutoModelForCausalLM, BitsAndBytesConfig

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"  # 예시

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,  # GPU가 bf16 지원이면 권장
)

tokenizer = AutoTokenizer.from_pretrained(model_id, use_fast=True)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="flash_attention_2",  # 안 되면 "sdpa"로
)

model.eval()

prompt = "로컬 LLM이 느릴 때 가장 먼저 점검할 것은?"
inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

with torch.inference_mode():
    out = model.generate(
        **inputs,
        max_new_tokens=128,
        do_sample=False,
        use_cache=True,
    )

print(tokenizer.decode(out[0], skip_special_tokens=True))

attn_implementation이 안 먹히는 경우

환경에 따라 flash_attention_2가 설치되어 있어야 하거나, 모델 아키텍처가 해당 경로를 지원해야 합니다. 실패 시 다음을 우선 시도하세요.

  • attn_implementation="sdpa"로 변경
  • PyTorch 버전 확인(최소 2.x 권장)
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    device_map="auto",
    torch_dtype=torch.bfloat16,
    attn_implementation="sdpa",
)

sdpa만으로도 기본 eager attention 대비 개선되는 경우가 많습니다.

속도 측정: tok/s만 보지 말고 prefill과 decode를 분리

체감 성능은 보통 다음 두 값으로 나뉩니다.

  • 첫 토큰까지 걸리는 시간(대화 UX에서 중요)
  • 생성 tok/s(길게 생성할 때 중요)

간단 측정 예시입니다.

import time
import torch

def benchmark_generate(model, tokenizer, prompt, max_new_tokens=128):
    inputs = tokenizer(prompt, return_tensors="pt").to(model.device)

    # 워밍업
    with torch.inference_mode():
        _ = model.generate(**inputs, max_new_tokens=8, do_sample=False, use_cache=True)

    torch.cuda.synchronize()
    t0 = time.time()

    with torch.inference_mode():
        out = model.generate(
            **inputs,
            max_new_tokens=max_new_tokens,
            do_sample=False,
            use_cache=True,
            return_dict_in_generate=True,
            output_scores=False,
        )

    torch.cuda.synchronize()
    t1 = time.time()

    total_new = out.sequences.shape[1] - inputs["input_ids"].shape[1]
    tok_s = total_new / (t1 - t0)
    return tok_s

print("tok/s:", benchmark_generate(model, tokenizer, "긴 프롬프트를 넣고 테스트해보세요.", 256))

정확히 prefill/decode를 분리하려면 generate 내부를 더 세밀하게 나눠야 하지만, 위 정도만 해도 “설정 변경이 성능에 영향이 있는지”는 충분히 비교할 수 있습니다.

4bit가 항상 빨라지지 않는 이유 (그래도 쓰는 이유)

4bit 양자화는 본질적으로 메모리 대역폭과 캐시 효율을 개선할 수 있지만, 동시에 다음 비용이 생길 수 있습니다.

  • 연산 중간에 dequantize가 필요
  • 커널 최적화가 GPU/드라이버/라이브러리 조합에 따라 달라짐

그럼에도 로컬 환경에서는 대개

  • 16bit로는 VRAM이 부족해 offload가 발생하고
  • offload로 인한 PCIe 병목이 훨씬 치명적

이기 때문에, 4bit로 “GPU 상주”를 유지하는 것만으로도 결과적으로 빨라지는 경우가 많습니다.

FlashAttention이 특히 효과적인 상황

다음 조건에서 개선 폭이 커집니다.

  • 컨텍스트 길이가 길다(예: 4k, 8k 이상)
  • 배치가 약간이라도 있다(여러 요청을 묶거나, speculative decoding 실험 등)
  • GPU가 충분히 강한데도 tok/s가 기대보다 낮다

반대로 프롬프트가 매우 짧고, 이미 decode만 하는 상황에서는 체감이 제한적일 수 있습니다.

자주 터지는 함정 체크리스트

use_cache가 꺼져 있다

커스텀 코드나 파이프라인에서 use_cache=False가 섞이면 속도가 급락합니다.

gen = model.generation_config
print("use_cache:", gen.use_cache)

# 필요 시 강제
model.generation_config.use_cache = True

dtype 혼합으로 인한 성능 저하

  • 모델 로드는 bf16인데 compute dtype이 fp16으로 섞이는 경우
  • CPU로 떨어지는 텐서가 생기는 경우

가능하면 다음처럼 일관성을 유지하세요.

bnb_config = BitsAndBytesConfig(
    load_in_4bit=True,
    bnb_4bit_quant_type="nf4",
    bnb_4bit_use_double_quant=True,
    bnb_4bit_compute_dtype=torch.bfloat16,
)

model = AutoModelForCausalLM.from_pretrained(
    model_id,
    quantization_config=bnb_config,
    torch_dtype=torch.bfloat16,
    device_map="auto",
)

프롬프트가 너무 길어 첫 응답이 느리다

“tok/s는 괜찮은데 첫 응답이 늦다”면 prefill 병목입니다. 해결은 보통 다음 중 하나입니다.

  • 컨텍스트를 줄이기(RAG면 top-k 축소, chunk 최적화)
  • 요약/메모리 전략 도입
  • FlashAttention/SDPA 적용

RAG 쪽에서 불필요한 문서가 섞여 컨텍스트가 커지는 문제는 인덱스/필터 설계가 원인인 경우도 있습니다. 벡터 검색을 쓴다면 Pinecone·Milvus 필터 누락을 막는 인덱스 설계도 함께 점검해보세요.

설치/환경 팁: “가속이 적용됐는지” 확인하는 요령

가속이 적용되지 않는 가장 흔한 원인은 “설치했는데 실제로 그 경로를 안 타는 것”입니다.

  • attn_implementation을 명시했는지 확인
  • 실행 로그/경고에서 fallback 메시지가 있는지 확인
  • 가능하면 torch.backends.cuda.matmul.allow_tf32 = True(Ampere 이상)도 검토
import torch

torch.backends.cuda.matmul.allow_tf32 = True
torch.set_float32_matmul_precision("high")

TF32는 정확도에 민감한 학습이 아니라 추론에서는 종종 좋은 트레이드오프가 됩니다.

결론: “4bit는 생존, FlashAttention은 가속”으로 접근하자

로컬 LLM이 느릴 때의 우선순위를 요약하면 다음과 같습니다.

  1. VRAM이 부족하면 bitsandbytes 4bit로 offload를 최소화한다.
  2. 그 다음 FlashAttention-2 또는 PyTorch SDPA로 어텐션 커널을 교체한다.
  3. use_cache=True, dtype 일관성, 프롬프트 길이(=prefill) 병목을 점검한다.
  4. tok/s만 보지 말고 첫 토큰 지연과 생성 속도를 분리해 측정한다.

이 조합만 제대로 맞춰도 “로컬 LLM은 원래 느리다”가 아니라, 설정과 커널 선택의 문제였다는 걸 체감하게 됩니다.

Discuss on Twitter • View on GitHub