[Bug]: KV cache quant - cannot be loaded in vllm #1949
Description
⚙️ Your current environment
The output of python collect_env.py
### Environment Information ###
Operating System: `Linux-6.17.1-arch1-1-x86_64-with-glibc2.42`
Python Version: `3.12.9 (main, Mar 17 2025, 21:01:58) [Clang 20.1.0 ]`
llm-compressor Version: `0.8.1`
compressed-tensors Version: `0.12.2`
transformers Version: `4.56.2`
torch Version: `2.8.0+cu129`
CUDA Devices: `['NVIDIA RTX PRO 6000 Blackwell Workstation Edition', 'NVIDIA RTX PRO 6000 Blackwell Workstation Edition']`
AMD Devices: `None`
vLLM version: vLLM API server version 0.11.1rc2.dev150+g5c2acb270.d20251018
🐛 Describe the bug
A follow-up to #1881
As mentioned in my last comment, upgrading from llmcompressor 0.7.1 -> 0.8.1 allowed me to complete my quantization script.
However the resulting model does not load in vllm with --kv-cache-dtype='fp8', as it triggers an assert
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] WorkerProc failed to start.
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] Traceback (most recent call last):
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/v1/executor/multiproc_executor.py", line 602, in worker_main
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] worker = WorkerProc(*args, **kwargs)
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] ^^^^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/v1/executor/multiproc_executor.py", line 457, in __init__
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] self.worker.load_model()
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/v1/worker/gpu_worker.py", line 230, in load_model
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] self.model_runner.load_model(eep_scale_up=eep_scale_up)
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/v1/worker/gpu_model_runner.py", line 2867, in load_model
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] self.model = model_loader.load_model(
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] ^^^^^^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/model_executor/model_loader/base_loader.py", line 56, in load_model
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] process_weights_after_loading(model, model_config, target_device)
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/model_executor/model_loader/utils.py", line 117, in process_weights_after_loading
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] quant_method.process_weights_after_loading(module)
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] File "/workspace/vllm/vllm/model_executor/layers/quantization/kv_cache.py", line 69, in process_weights_after_loading
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] assert layer.k_scale > 0.0
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] ^^^^^^^^^^^^^^^^^^^
(Worker_TP0 pid=224) ERROR 10-19 14:56:04 [multiproc_executor.py:628] AssertionError
🛠️ Steps to reproduce
The script is there https://github.com/mratsim/quantizers/blob/b97bdd8/main_seed-oss-fp8-kv8.py
from datasets import load_dataset
from transformers import AutoModelForCausalLM, AutoTokenizer
from llmcompressor import oneshot
from llmcompressor.utils import dispatch_for_generation
from llmcompressor.modifiers.quantization import QuantizationModifier
from compressed_tensors.quantization import (
QuantizationArgs,
QuantizationScheme,
QuantizationStrategy,
QuantizationType,
)
CALIBRATION_DATASET="HuggingFaceH4/ultrachat_200k"
DATASET_SPLIT="train_sft"
SHUFFLE_SEED=42
NUM_CALIBRATION_SAMPLES = 512
MAX_SEQUENCE_LENGTH = 4096
MODEL_ID = "ByteDance-Seed/Seed-OSS-36B-Instruct"
MODEL_OUT = MODEL_ID.split("/")[1] + "-FP8-KV8"
# Load model.
model = AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtype="auto")
tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model.generation_config.do_sample=True
# Dataset processing
ds = load_dataset(CALIBRATION_DATASET, split=f"{DATASET_SPLIT}[:{NUM_CALIBRATION_SAMPLES}]")
ds = ds.shuffle(seed=42)
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)
recipe = [
QuantizationModifier(
ignore=["lm_head"],
# DeepSeek V3 style block quantization + dynamic per token quantization
config_groups={
"group_0": QuantizationScheme(
targets=["Linear"],
weights=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
dynamic=False,
symmetric=True,
strategy=QuantizationStrategy.BLOCK,
block_structure=[128, 128],
),
input_activations=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
strategy=QuantizationStrategy.GROUP,
symmetric=True,
dynamic=True,
observer=None,
group_size=128,
),
),
},
kv_cache_scheme=QuantizationArgs(
num_bits=8,
type=QuantizationType.FLOAT,
dynamic=False,
symmetric=True,
strategy=QuantizationStrategy.TENSOR,
),
)
]
oneshot(
# pipeline="basic",
model=model,
recipe=recipe,
dataset=ds,
max_seq_length=MAX_SEQUENCE_LENGTH,
num_calibration_samples=NUM_CALIBRATION_SAMPLES,
)
# Save to disk in compressed-tensors format.
model.save_pretrained(MODEL_OUT, save_compressed=True)
tokenizer.save_pretrained(MODEL_OUT)
print(f'SUCCESS: files saved in {MODEL_OUT}')