양자화된 KV 캐시

양자화된 KV 캐시 (Quantized KV Cache)

FP8 KV 캐시로 모델의 Key-Value 캐시를 양자화하면 메모리 풋프린트를 크게 줄일 수 있어요. 이 최적화로 더 많은 토큰을 메모리에 저장할 수 있어 처리량이 개선되고 더 긴 컨텍스트 윈도우를 지원하게 됩니다.

참고: Flash Attention 3 백엔드와 FP8 KV 캐시를 함께 사용하면 어텐션 연산도 양자화된(FP8) 도메인에서 수행됩니다. 이 구성에서는 key와 value뿐 아니라 쿼리도 FP8로 양자화됩니다.

출처: 문서

본문

지원되는 FP8 KV-캐시 양자화 체계 (Supported FP8 KV-Cache Quantization Schemes)

vLLM은 FP8 KV-캐시에 대해 두 가지 주요 양자화 전략을 지원합니다.

Per-tensor 양자화: 각 Q, K, V 텐서에 단일 스케일이 적용됩니다. (q/k/v_scale = [1])
Per-attention-head 양자화: 각 스케일이 어텐션 헤드에 해당합니다. q_scale = [num_heads], k/v_scale = [num_kv_heads].

참고: Per-attention-head 양자화는 현재 Flash Attention 백엔드에서만 사용할 수 있으며 llm-compressor 가 제공하는 보정 경로가 필요합니다.

스케일 보정 접근법 (Scale Calibration Approaches)

vLLM에서 양자화 스케일을 계산하는 방식을 세 가지로 구성할 수 있습니다.

보정 없음(기본 스케일): 모든 양자화 스케일이 1.0 으로 설정됩니다. 구성: kv_cache_dtype="fp8".
[권장] 데이터셋으로 보정(llm-compressor 경유): 스케일을 큐레이트된 보정 데이터셋으로 추정해 최대 정확도를 얻습니다. 이는 llm-compressor 라이브러리가 필요합니다. 아래 예시를 참고하세요!

추가 `kv_cache_dtype` 옵션

kv_cache_dtype="auto": 모델의 기본 데이터 타입을 사용합니다.
kv_cache_dtype="fp8_e4m3": CUDA 11.8+ 및 ROCm(AMD GPU)에서 지원됩니다.
kv_cache_dtype="fp8_e5m2": CUDA 11.8+에서 지원됩니다.

KV-캐시 양자화에서 특정 레이어 건너뛰기 (Skipping Specific Layers from KV-Cache Quantization)

일부 어텐션 레이어 유형(예: 슬라이딩 윈도우)은 KV-캐시 양자화에 더 민감합니다. --kv-cache-dtype-skip-layers 플래그는 나머지 레이어를 선택한 양자화 dtype 아래로 유지하면서 지정된 레이어를 모델의 네이티브 dtype으로 남깁니다. 이 플래그는 레이어 인덱스나 레이어 타입 이름을 받습니다.

# Skip every sliding-window attention layer.
vllm serve <model> \
  --kv-cache-dtype fp8 \
  --kv-cache-dtype-skip-layers sliding_window

# Skip specific layer indices.
vllm serve <model> \
  --kv-cache-dtype fp8 \
  --kv-cache-dtype-skip-layers 0 1 23

프로그래매틱 사용법:

llm = LLM(
    model="meta-llama/Llama-3.1-8B-Instruct",
    kv_cache_dtype="fp8",
    kv_cache_dtype_skip_layers=["sliding_window"],
)

예시 (Examples)

1. 보정 없음 (`kv_cache_dtype="fp8"`)

모든 양자화 스케일이 1.0으로 설정됩니다.

from vllm import LLM, SamplingParams

sampling_params = SamplingParams(temperature=0.7, top_p=0.8)
llm = LLM(
    model="meta-llama/Llama-2-7b-chat-hf",
    kv_cache_dtype="fp8",
)
prompt = "London is the capital of"
out = llm.generate(prompt, sampling_params)[0].outputs[0].text
print(out)

2. [권장] 데이터셋으로 보정 (`llm-compressor` 사용)

최고 품질의 양자화를 위해 llm-compressor 로 데이터셋에 대해 보정하는 것을 권장합니다. 이는 per-attention-head 양자화 같은 고급 전략을 가능하게 합니다.

필요한 패키지를 설치합니다.

pip install llmcompressor

예시: Llama Attention & KV Cache를 FP8로 양자화.

"""
Quantize Llama attention + KV cache to FP8 (choose either 'tensor' or 'attn_head' strategy)
using llm-compressor one-shot calibration.
"""

from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer

from llmcompressor import oneshot
from llmcompressor.modifiers.quantization import QuantizationModifier
from compressed_tensors.quantization import QuantizationScheme, QuantizationArgs

# -----------------------------
# Config
# -----------------------------
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
STRATEGY = "tensor"       # or "attn_head"
NUM_CALIB_SAMPLES = 512   # Good starting value
MAX_SEQ_LEN = 2048

# -----------------------------
# Helpers
# -----------------------------
def process_and_tokenize(example, tokenizer: AutoTokenizer):
    """Convert chat messages to tokens."""
    text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
    return tokenizer(
        text,
        padding=False,
        max_length=MAX_SEQ_LEN,
        truncation=True,
        add_special_tokens=False,
    )

def build_recipe(strategy: str) -> QuantizationModifier:
    fp8_args = QuantizationArgs(num_bits=8, type="float", strategy=strategy)
    return QuantizationModifier(
        config_groups={
            "attention": QuantizationScheme(
                targets=["LlamaAttention"],  # Quantize queries: q_scale
                input_activations=fp8_args,
            )
        },
        kv_cache_scheme=fp8_args,           # Quantize KV cache: k/v_scale
    )

# -----------------------------
# Main
# -----------------------------
def main():
    model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
    tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
    ds = load_dataset(DATASET_ID, split=f"{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]")
    ds = ds.shuffle(seed=42)
    ds = ds.map(
        lambda ex: process_and_tokenize(ex, tokenizer),
        remove_columns=ds.column_names,
    )

    recipe = build_recipe(STRATEGY)
    oneshot(
        model=model,
        dataset=ds,
        recipe=recipe,
        max_seq_length=MAX_SEQ_LEN,
        num_calibration_samples=NUM_CALIB_SAMPLES,
    )

    save_dir = f"{MODEL_ID.rstrip('/').split('/')[-1]}-kvattn-fp8-{STRATEGY}"
    model.save_pretrained(save_dir, save_compressed=True)
    tokenizer.save_pretrained(save_dir)

if __name__ == "__main__":
    main()

더 상세하고 최신 예시는 llm-compressor 공식 예시 를 참고하세요.

양자화된 KV 캐시