量化 KV 快取¶
FP8 KV 快取¶
將 KV 快取量化為 FP8 可減小其記憶體佔用。這增加了可儲存在快取中的 token 數量，從而提高了吞吐量。
FP8 格式¶
OCP (Open Compute Project) 規定了兩種常見的 8 位浮點資料格式
- E5M2 (5 位指數和 2 位尾數)
- E4M3FN (4 位指數和 3 位尾數，通常簡稱為 E4M3)
與 E5M2 相比，E4M3 格式提供了更高的精度。但是，由於其較小的動態範圍（±240.0），E4M3 通常需要在每個量化張量旁邊使用更高精度的 (FP32) 縮放因子。
當前限制¶
目前，僅支援每張量（標量）縮放因子。正在開發中以支援更細粒度的縮放因子（例如每通道）。
FP8 KV 快取的工作原理¶
FP8 KV 快取的實現遵循以下工作流程
- 儲存：鍵和值張量使用縮放因子量化為 FP8 格式，然後儲存在 KV 快取中
- 檢索：當需要進行注意力計算時，快取的 KV 張量會被反量化回更高精度 (FP16/BF16)
- 注意力：注意力-值乘法（softmax 輸出 × V）使用反量化的高精度 V 張量進行計算
這意味著最終的注意力計算是在反量化值上進行的，而不是 FP8 張量。量化減少了儲存期間的記憶體使用，但透過在實際注意力操作中使用更高精度來保持計算準確性。
效能影響¶
當前的 FP8 KV 快取實現主要透過允許大約兩倍的 KV 快取分配空間來提高吞吐量。這可以實現
- 為單個請求處理更長的上下文長度，或
- 處理更多併發請求批次
但是，目前沒有延遲改進，因為該實現尚未包含融合的反量化和注意力操作。未來的版本將支援具有硬體加速的量化注意力，這將提供額外的效能優勢。雖然最新的晶片產品（例如 AMD MI300、NVIDIA Hopper 或更高版本）支援 FP8 與其他格式（fp32、fp16、bf16）之間的本機硬體轉換，但尚未完全實現這一優勢。
研究表明，FP8 E4M3 量化通常只會輕微降低推理精度，使其成為吞吐量最佳化的實用選擇。
使用示例¶
以下是如何啟用 FP8 量化的示例
程式碼
# To calculate kv cache scales on the fly enable the calculate_kv_scales
# parameter
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",
calculate_kv_scales=True,
)
prompt = "London is the capital of"
out = llm.generate(prompt, sampling_params)[0].outputs[0].text
print(out)
kv_cache_dtype 引數指定 KV 快取儲存的資料型別
"auto"：使用模型的預設“未量化”資料型別"fp8"或"fp8_e4m3"：支援 CUDA 11.8+ 和 ROCm (AMD GPU)"fp8_e5m2"：支援 CUDA 11.8+
經過校準的縮放因子以提高精度¶
為了在使用 FP8 KV 快取時獲得最佳的模型質量，我們建議使用針對代表性推理資料進行調整的校準縮放因子。LLM Compressor 是此過程的推薦工具。
安裝¶
首先，安裝所需的依賴項
示例用法¶
以下是一個使用 meta-llama/Llama-3.1-8B-Instruct 的完整示例（大多數模型可以使用相同的模式）
程式碼
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
# Select model and load it
MODEL_ID = "meta-llama/Llama-3.1-8B-Instruct"
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, device_map="auto", dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
# Select calibration dataset
DATASET_ID = "HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT = "train_sft"
# Configure calibration parameters
NUM_CALIBRATION_SAMPLES = 512 # 512 samples is a good starting point
MAX_SEQUENCE_LENGTH = 2048
# Load and preprocess dataset
ds = load_dataset(DATASET_ID, split=DATASET_SPLIT)
ds = ds.shuffle(seed=42).select(range(NUM_CALIBRATION_SAMPLES))
def process_and_tokenize(example):
text = tokenizer.apply_chat_template(example["messages"], tokenize=False)
return tokenizer(
text,
padding=False,
max_length=MAX_SEQUENCE_LENGTH,
truncation=True,
add_special_tokens=False,
)
ds = ds.map(process_and_tokenize, remove_columns=ds.column_names)
# Configure quantization settings
recipe = """
quant_stage:
quant_modifiers:
QuantizationModifier:
kv_cache_scheme:
num_bits: 8
type: float
strategy: tensor
dynamic: false
symmetric: true
"""
# Apply quantization
oneshot(
model=model,
dataset=ds,
recipe=recipe,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)
# Save quantized model: Llama-3.1-8B-Instruct-FP8-KV
SAVE_DIR = MODEL_ID.split("/")[1] + "-FP8-KV"
model.save_pretrained(SAVE_DIR, save_compressed=True)
tokenizer.save_pretrained(SAVE_DIR)
上面的指令碼將在您當前的目錄中建立一個資料夾，其中包含您的量化模型（例如，Llama-3.1-8B-Instruct-FP8-KV）以及校準的縮放因子。
執行模型時，您必須指定 kv_cache_dtype="fp8" 才能啟用 KV 快取量化並使用縮放因子。