量化模型AWQ加载时显存不足验证KV cache量化开关影响

当AWQ量化模型在加载阶段就报显存不足，KV cache量化开关往往不是第一排查点。KV cache分配通常在模型权重加载完成后的预热或生成阶段发生，因此在“从预训练模型加载权重”这一步出现CUDA out of memory，多半与KV cache量化无关。需要先定位显存不足发生在哪个阶段，再决定是否调整KV cache量化开关。

KV cache量化开关影响的是推理时KV缓存的内存占用，不改变模型权重加载阶段所需的显存。如果显存不足发生在from_pretrained阶段，应先检查权重精度、设备映射和临时缓冲区；如果显存不足发生在生成token或长序列推理阶段，才需要验证KV cache量化开关的实际收益。正确的做法是分阶段对照，不要一上来就改量化开关。

先弄清KV cache量化开关的作用域

AWQ量化压缩的是模型权重，也就是在加载模型时，权重矩阵从int4/int8等形式进入显存，占用的体积由权重位数和模型结构决定。KV cache是推理过程中逐渐累积的key/value张量，每个生成的token都会追加一部分。KV cache量化开关通常用于改变这一部分缓存的数据类型，比如从fp16压到fp8或int8，从而降低每个token占用的显存。

这个开关只影响“生成阶段”的显存增长，不改变模型权重从磁盘加载到显存时的总量。如果您的场景是“模型能加载，但生成几个token后OOM”，那么开启KV cache量化值得验证。如果“模型本身加载不进去”，那与KV cache量化没有直接关系。

先定位显存不足发生在哪一步

建议先用最小脚本把加载和生成拆开，观察各自阶段的显存峰值。下面是一个基于Transformers的通用骨架，用来检查权重加载阶段的显存占用：

# 阶段1：只加载权重，不生成任何token
import torch
from transformers import AutoModelForCausalLM, AutoTokenizer

model_path = "/path/to/your-awq-model"
tokenizer = AutoTokenizer.from_pretrained(model_path)
model = AutoModelForCausalLM.from_pretrained(model_path, device_map="cuda")

print("权重加载后已分配显存:", torch.cuda.memory_allocated())
print("权重加载后峰值显存:", torch.cuda.max_memory_allocated())

如果这段代码执行到from_pretrained时就报OOM，说明瓶颈在权重本身、CUDA上下文或设备映射。此时KV cache量化开关无法解决。可尝试的动作包括：换用更小的max_memory分片、增加GPU数量、使用CPU offload、确认是否有多进程各占一份权重导致显存翻倍，以及检查AWQ权重文件是否完整加载了不需要的嵌入层。

在推理阶段对照KV cache量化开关

如果模型能加载，但生成时OOM，才需要对比KV cache量化开关。以vLLM这类常见推理框架为例，启动参数中通常会有一个类似`--kv-cache-dtype`的选项，支持auto、fp8_e5m2等取值。可以固定相同的模型路径、最大长度和batch大小，分别用不开启该选项和开启该选项启动服务，再发送请求观察显存变化。

# 对照组：不开启KV cache量化
python -m vllm.entrypoints.openai.api_server \
    `--model` /path/to/your-awq-model \
    `--max-model-len` 4096 \
    `--gpu-memory-utilization` 0.9

# 实验组：开启KV cache量化（具体参数名需要以当前框架版本为准）
python -m vllm.entrypoints.openai.api_server \
    `--model` /path/to/your-awq-model \
    `--max-model-len` 4096 \
    `--kv-cache-dtype` fp8_e5m2 \
    `--gpu-memory-utilization` 0.9

如果使用Transformers的生成接口，可以观察生成固定长度序列时的显存曲线：

# 阶段2：固定输入长度，生成固定长度的token
inputs = tokenizer("测试KV cache显存占用，请输出一段内容。", return_tensors="pt").to("cuda")
with torch.no_grad():
    generated = model.generate(**inputs, max_new_tokens=512)

print("生成后已分配显存:", torch.cuda.memory_allocated())
print("生成后峰值显存:", torch.cuda.max_memory_allocated())

对比两组代码在“权重加载后”和“生成后”的显存差值，差值主要就是KV cache增长量。如果开关开启后差值明显变小，同时没有触发OOM，则说明KV cache量化开关对当前场景有效。

验证时应控制哪些变量

为了得到可复现的对照结果，建议固定以下变量：模型路径、最大序列长度、batch size、生成token数量、输入prompt长度、并发请求数、GPU驱动和框架版本。每次只改变KV cache量化开关，并记录以下信息：

权重加载阶段峰值显存：用于排除权重问题。
生成阶段峰值显存：用于观察KV cache量化实际效果。
是否触发OOM：最直接的判断依据。
输出质量：KV cache量化可能带来精度损失，需要主观或自动指标评估。

一个值得注意的边界：KV cache量化开关通常只在长上下文或大batch场景下才有意义。如果模型本身仅支持短对话，并且显存不足发生在单次短生成时，那么问题可能出在权重加载缓存或CUDA上下文，优先检查CUDA_MEMORY_FRAGMENTATION、页面大小和进程占用的非模型显存，而不是继续调KV cache量化。

最后建议，把“加载显存不足”和“推理显存不足”分开记录。若您的显存不足仅在加载阶段出现，先处理权重加载路径，KV cache量化开关可以暂时关闭，避免引入不必要的精度变化。若两种情况都出现，则先解决权重加载，再逐步开启KV cache量化并验证长序列稳定性。