vLLM FP8 量化 KV Cache 完整指南格式选择、Scale 校准与 llm-compressor 实操【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllmKV Cache 是长上下文推理中 GPU 显存的主要消耗来源之一。把 Key/Value 缓存量化为 FP8可以将缓存显存占用减半左右让同一块 GPU 装下更多 token从而直接提升吞吐并支撑更长的上下文窗口。本文以 vLLM 官方量化文档 quantized_kvcache.md 为主体结合 CacheConfig、EngineArgs 参数定义、Attention 层实现 与 FlashAttention 后端 等源码系统讲解 vLLM FP8 KV Cache 的量化粒度、存储格式、三种 scale 获取方式、按层跳过量化等机制并给出可直接运行的 Python 与命令行示例。FP8 量化 KV Cache 的价值与适用场景LLM 推理过程中KV Cache 的大小随并发序列数、上下文长度线性增长往往是显存中最大的动态占用项。将 K/V 张量以 FP88 位浮点相对 BF16/FP16 减半存储可以在几乎不改变计算图的前提下显著缩小缓存体积。收益主要体现在两方面同卡可容纳更多 tokenKV Cache 更小 → 可分配的 cache block 更多 → 单批可承载的有效 token 数更大支持更长上下文窗口在gpu_memory_utilization约束下更紧凑的缓存让长序列请求更容易获得完整可用的 KV 空间。量化是有代价的精度下降可能带来端到端准确率损失尤其是当 scale缩放系数不合适时。因此 vLLM 在 CacheConfig 中会以日志明确提示用户使用量化 KV Cache dtype 会降低显存占用并提升性能但 may cause accuracy drop without a proper scaling factor缺少合适缩放因子时可能产生精度损失。注意FlashAttention 3 后端当使用 FA3 后端 FP8 KV Cache 时注意力计算本身也在量化FP8域中进行此时 Query 也会和 Key/Value 一样被量化为 FP8。这是文档中特别提示的一个后端行为差异。kv_cache_dtypeFP8 缓存存储格式怎么选KV Cache 的存储类型由kv_cache_dtype控制。核心取值定义在 CacheConfig.CacheDType与 FP8 相关的选项如下取值含义平台支持auto使用模型默认数据类型等价于不做显式量化通用fp8短别名等价于fp8_e4m3取决于平台支持fp8_e4m3FP8 E4M3 格式精度更高推荐首选CUDA 11.8、ROCmAMD GPUfp8_e5m2FP8 E5M2 格式动态范围更大、精度较低CUDA 11.8从 torch_utils.py 可以看到别名映射关系fp8 - fp8_e4m3因此直接写fp8与写fp8_e4m3等价。在 vLLM 内部所有 fp8 前缀的缓存 dtype 都会被 is_quantized_kv_cache 识别为量化缓存在部分注意力量化实现里fp8*dtype 对应的底层张量类型为torch.uint8见 STR_DTYPE_TO_TORCH_DTYPE实际以 FP8 语义解释位模式。需要说明的是E4M3 与 E5M2 的可用性取决于后端实现与硬件算子。除上述选项外当前仓库的 CacheDType 还枚举了fp8_incIntel Gaudi、fp8_ds_mlaDeepSeek MLA 相关以及fp8_per_token_head、int4_per_token_head等每 token-头动态 scale 类型以及 NVFP4 系列nvfp4、nvfp4_4over6可视为同一配置体系下的扩展面本文聚焦 FP8 主路径。CLI 与编程接口两种配置方式CLI 服务端使用缓存参数组中定义的--kv-cache-dtype见 arg_utils.py 参数组注册vllm serve model --kv-cache-dtype fp8离线/编程接口在构造LLM或AsyncLLM时传入同名关键字二者底层最终汇入同一个CacheConfigfrom vllm import LLM llm LLM( modelmeta-llama/Llama-3.1-8B-Instruct, kv_cache_dtypefp8, # 等价于 fp8_e4m3 )当kv_cache_dtypeauto时引擎启动阶段会通过resolve_kv_cache_dtype_stringarg_utils.py结合模型配置把 auto 解析为实际 dtype 后写入CacheConfig.cache_dtype。同时若模型 checkpoint 中自带 FP8 KV Cache 量化声明见下文 llm-compressor 校准且用户没有显式指定 dtypeattention.py 会自动把缓存 dtype 解析为fp8保证校准过就生效。两种 FP8 量化粒度Per-Tensor 与 Per-Attention-Head对同一个 FP8 缓存vLLM 支持两种量化/反量化 scale 粒度方案scale 形状说明Per-tensor 量化q/k/v_scale [1]Q、K、V 各自整体共享一个标量 scale开销最小精度相对粗Per-attention-head 量化q_scale [num_heads]k/v_scale [num_kv_heads]每个注意力头独立 scale精度更高需要保存/计算的头级 scale 张量更多其中 Per-attention-head 方案当前仅 FlashAttention 后端支持需要由llm-compressor提供的校准路径预先产出各头 scale无法在 vLLM 运行时凭空推断。从源码看Attention 层在初始化时会读取量化配置中的kv_cache_scheme当strategy attn_head时置use_per_head_quant_scales Trueattention.py后续将头级 scale 张量交由注意力后端使用在 FlashAttention 后端 中可以见到对layer._q_scale、layer._k_scale等做 shape 扩展expand(descale_shape)以匹配每头解量化的实现痕迹这正是每头 scale 落到 kernel 层的通路。Scale 校准的两种配置途径以及第三种隐式途径选哪个 dtype只是第一步更关键的问题是量化 scale 从哪里来。scale 若全为 1.0本质是截断式量化误差最大scale 越贴合真实激活分布精度越好。vLLM 提供以下方式方式一不做校准默认 scale 1.0仅通过kv_cache_dtypefp8开启量化所有 scale 一律取1.0llm LLM(modelmeta-llama/Llama-2-7b-chat-hf, kv_cache_dtypefp8)这种方式部署最简单、零额外依赖适合快速验证显存/吞吐收益但由于缺少与真实激活匹配的 scale精度损失通常最明显。引擎也会在 CacheConfig 校验 时打印对应提示。方式二用校准数据集估算 scale推荐对追求质量的生产场景vLLM 官方推荐使用llm-compressorvLLM 生态下的模型压缩库在一次性one-shot校准过程中用小规模有代表性的数据集统计出 Q/K/V 激活分布从而估计出最优 scale。它同时是开启 per-attention-head 高级量化的前提。具体操作见本文后面的端到端实操示例。第三种隐式途径直接消费 checkpoint 里固化好的 scalellm-compressor 校准完成后会把量化方案含kv_cache_scheme与 scale 系数随模型一起保存save_pretrained(..., save_compressedTrue)。之后在 vLLM 中加载该模型时如果用户没有显式传kv_cache_dtype即autoattention.py 会读取quant_config.kv_cache_scheme并自动切换到fp8缓存若 scheme 的 strategy 为attn_head则启用 per-head scale 通路同一文件 L288-L292。这意味着校准工作只需要做一次产出的模型即可在后续部署中自动复用 scale无需每次启动重复校准。对量化敏感层单独豁免--kv-cache-dtype-skip-layers并非所有注意力层都适合量化。例如sliding-window 注意力层常被用于长序列局部建模对 KV Cache 量化更敏感量化后质量下降更明显。kv_cache_dtype_skip_layers允许你把指定层留在模型原生 dtype而其余层继续使用所选量化 dtype实现量化与精度的折中。该参数在 CacheConfig 中定义为字符串列表接受两类输入注意力类型名如sliding_window表示匹配所有该类层层序号如0、1、23表示只豁免指定索引的层。命令行用法# 跳过所有 sliding-window 注意力层 vllm serve model \ --kv-cache-dtype fp8 \ --kv-cache-dtype-skip-layers sliding_window # 跳过第 0、1、23 层按层索引 vllm serve model \ --kv-cache-dtype fp8 \ --kv-cache-dtype-skip-layers 0 1 23编程接口用法from vllm import LLM llm LLM( modelmeta-llama/Llama-3.1-8B-Instruct, kv_cache_dtypefp8, kv_cache_dtype_skip_layers[sliding_window], )底层实现机制这个功能的判定逻辑位于 Attention 层初始化其执行流程为若该层存在 sliding-windowsliding_window is not None且 skip 列表包含sliding_window→ 命中否则用模块前缀名如model.layers.0.self_attn经 extract_layer_index 提取层索引若该索引以字符串形式出现在 skip 列表 → 命中命中后该层的kv_cache_dtype被改写为auto即回退到模型原生 dtype其他层不受影响。从 CacheConfig 的注释还能看到配套机制跳过层仍要参与 block-size 对齐必要时通过skip_page_size_padded把未量化层补齐到量化主层 page 大小保证混合缓存布局在页级可对齐。测试用例如tests/kernels/attention/test_attention.py对kv_cache_dtype相关矩阵的覆盖会验证 skip 列表与不同 dtype 组合的行为一致性。实操示例一无校准直接开启 FP8最简上手路径将示例模型改为你环境中实际可下载的模型即可仓库内示例使用 Llama-2-7B-chatfrom vllm import LLM, SamplingParams sampling_params SamplingParams(temperature0.7, top_p0.8) llm LLM( modelmeta-llama/Llama-2-7b-chat-hf, kv_cache_dtypefp8, ) prompt London is the capital of out llm.generate(prompt, sampling_params)[0].outputs[0].text print(out)运行后对比同模型auto与fp8两档fp8档的 KV Cache 占用显著下降可观测的峰值显存随之降低同时由于所有 scale 均为 1.0输出质量可能略降——这正是需要方式二校准的理由。实操示例二llm-compressor 数据集校准推荐高质量量化路径。以下脚本用HuggingFaceH4/ultrachat_200k的少量样本对 Llama 3.1 8B 做一次性校准将注意力激活Q与 KV Cache同时量化为 FP8并支持在tensor与attn_head两种 strategy 间切换。安装依赖pip install llmcompressorllm-compressor 还会依赖datasets、transformers等示例脚本 import 中使用如缺失请一并安装。校准脚本 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], tokenizeFalse) return tokenizer( text, paddingFalse, max_lengthMAX_SEQ_LEN, truncationTrue, add_special_tokensFalse, ) def build_recipe(strategy: str) - QuantizationModifier: fp8_args QuantizationArgs(num_bits8, typefloat, strategystrategy) return QuantizationModifier( config_groups{ attention: QuantizationScheme( targets[LlamaAttention], # Quantize queries: q_scale input_activationsfp8_args, ) }, kv_cache_schemefp8_args, # Quantize KV cache: k/v_scale ) # ----------------------------- # Main # ----------------------------- def main(): model AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtypeauto) tokenizer AutoTokenizer.from_pretrained(MODEL_ID) ds load_dataset(DATASET_ID, splitf{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]) ds ds.shuffle(seed42) ds ds.map( lambda ex: process_and_tokenize(ex, tokenizer), remove_columnsds.column_names, ) recipe build_recipe(STRATEGY) oneshot( modelmodel, datasetds, reciperecipe, max_seq_lengthMAX_SEQ_LEN, num_calibration_samplesNUM_CALIB_SAMPLES, ) save_dir f{MODEL_ID.rstrip(/).split(/)[-1]}-kvattn-fp8-{STRATEGY} model.save_pretrained(save_dir, save_compressedTrue) tokenizer.save_pretrained(save_dir) if __name__ __main__: main()脚本要点拆解QuantizationModifierconfig_groups把LlamaAttention的输入激活即 Q按fp8_args量化为 FP8产出 per-tensor 或 per-head 的q_scalekv_cache_schemefp8_args对 KV Cache 应用同一 FP8 量化参数产出k/v_scalestrategytensorvsattn_head前者每张 Q/K/V 张量一个 scale后者为每个头一个 scale量化粒度更细、精度更好但仅在 vLLM 的 FlashAttention 后端下被消费支持save_compressedTrue把量化配置含kv_cache_scheme与权重一起写入 checkpoint即上文第三种隐式途径所依赖的载体校准样本数与序列长度是质量-耗时权衡点NUM_CALIB_SAMPLES从 512 起步通常是不错的起点MAX_SEQ_LEN应贴近你的真实部署长度。在 vLLM 中加载校准产物校准并保存后即可在 vLLM 中直接加载该目录无需再次指定 FP8from vllm import LLM llm LLM(modelLlama-3.1-8B-Instruct-kvattn-fp8-tensor)引擎会自动识别 checkpoint 中固化的 FP8 KV Cache 方案并启用对应 attention.py 的 auto 解析路径。若要显式覆盖例如退回 bfloat16 验证基线则显式传kv_cache_dtypebfloat16此时用户显式取值优先于 checkpoint 声明。提示llm-compressor 的量化 KV Cache 示例会随其版本持续演进建议以其仓库的examples/quantization_kv_cache目录为最新参考本仓库仅说明 vLLM 侧的消费与配置语义。落到引擎内部缓存 dtype 是如何贯穿配置到 kernel 的把上文涉及的各条线索串起来一次 FP8 KV Cache 请求的完整链路如下参数层用户经LLM(...)或vllm serve --kv-cache-dtype ...传入kv_cache_dtype与kv_cache_dtype_skip_layers引擎在 create_engine_config 中将二者解析/透传给CacheConfig配置层CacheConfigcache.py持有cache_dtype在 validators 中依据 is_quantized_kv_cache 打印显存降低 可能精度损失的运行提示模型层每个 Attention 层初始化时按 skip 规则决定自身 dtype并结合 checkpoint 的kv_cache_scheme决定是否启用 per-head scaleattention.py后端层注意力后端按 dtype 选择 kernel 路径FP8 路径下 FlashAttention 类实现会把key_cache/value_cache以 FP8 语义读取并利用_q_scale/_k_scale做反量化对齐flash_attn.py。这一层抽象意味着用户在配置层看到的kv_cache_dtype最终是由平台能力CUDA/ROCm/HPU 各自支持的具体格式、Attention 后端能力是否支持 per-head scale与模型/校准产物共同决定的。选型与注意事项小结首次尝试用kv_cache_dtypefp8无校准先量级评估显存与吞吐收益同时用原模型做精度对照追求质量用 llm-compressor 数据集校准优先strategytensor需要更高精度时再切attn_head注意其 FlashAttention 后端限定与模型结构限定示例仅覆盖 Llama 类架构的LlamaAttention混合模型/长上下文若模型含 sliding-window 层用--kv-cache-dtype-skip-layers sliding_window或指定层索引做定点豁免平台前提fp8_e4m3需 CUDA 11.8 或 ROCmfp8_e5m2需 CUDA 11.8auto始终可用精度兜底量化缓存本质上是以精度换显存/吞吐务必在真实业务样本上回归评测量化不总是无损。更完整、随版本更新的校准示例请以 llm-compressor 官方示例为准vLLM 侧的完整参数语义可继续研读 CacheConfig、EngineArgs 与 Attention 层相关源码以及仓库内 test_attention.py、test_cache.py 等针对缓存 dtype 与 kernel 的测试用例。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
vLLM FP8 量化 KV Cache 完整指南：格式选择、Scale 校准与 llm-compressor 实操
vLLM FP8 量化 KV Cache 完整指南格式选择、Scale 校准与 llm-compressor 实操【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllmKV Cache 是长上下文推理中 GPU 显存的主要消耗来源之一。把 Key/Value 缓存量化为 FP8可以将缓存显存占用减半左右让同一块 GPU 装下更多 token从而直接提升吞吐并支撑更长的上下文窗口。本文以 vLLM 官方量化文档 quantized_kvcache.md 为主体结合 CacheConfig、EngineArgs 参数定义、Attention 层实现 与 FlashAttention 后端 等源码系统讲解 vLLM FP8 KV Cache 的量化粒度、存储格式、三种 scale 获取方式、按层跳过量化等机制并给出可直接运行的 Python 与命令行示例。FP8 量化 KV Cache 的价值与适用场景LLM 推理过程中KV Cache 的大小随并发序列数、上下文长度线性增长往往是显存中最大的动态占用项。将 K/V 张量以 FP88 位浮点相对 BF16/FP16 减半存储可以在几乎不改变计算图的前提下显著缩小缓存体积。收益主要体现在两方面同卡可容纳更多 tokenKV Cache 更小 → 可分配的 cache block 更多 → 单批可承载的有效 token 数更大支持更长上下文窗口在gpu_memory_utilization约束下更紧凑的缓存让长序列请求更容易获得完整可用的 KV 空间。量化是有代价的精度下降可能带来端到端准确率损失尤其是当 scale缩放系数不合适时。因此 vLLM 在 CacheConfig 中会以日志明确提示用户使用量化 KV Cache dtype 会降低显存占用并提升性能但 may cause accuracy drop without a proper scaling factor缺少合适缩放因子时可能产生精度损失。注意FlashAttention 3 后端当使用 FA3 后端 FP8 KV Cache 时注意力计算本身也在量化FP8域中进行此时 Query 也会和 Key/Value 一样被量化为 FP8。这是文档中特别提示的一个后端行为差异。kv_cache_dtypeFP8 缓存存储格式怎么选KV Cache 的存储类型由kv_cache_dtype控制。核心取值定义在 CacheConfig.CacheDType与 FP8 相关的选项如下取值含义平台支持auto使用模型默认数据类型等价于不做显式量化通用fp8短别名等价于fp8_e4m3取决于平台支持fp8_e4m3FP8 E4M3 格式精度更高推荐首选CUDA 11.8、ROCmAMD GPUfp8_e5m2FP8 E5M2 格式动态范围更大、精度较低CUDA 11.8从 torch_utils.py 可以看到别名映射关系fp8 - fp8_e4m3因此直接写fp8与写fp8_e4m3等价。在 vLLM 内部所有 fp8 前缀的缓存 dtype 都会被 is_quantized_kv_cache 识别为量化缓存在部分注意力量化实现里fp8*dtype 对应的底层张量类型为torch.uint8见 STR_DTYPE_TO_TORCH_DTYPE实际以 FP8 语义解释位模式。需要说明的是E4M3 与 E5M2 的可用性取决于后端实现与硬件算子。除上述选项外当前仓库的 CacheDType 还枚举了fp8_incIntel Gaudi、fp8_ds_mlaDeepSeek MLA 相关以及fp8_per_token_head、int4_per_token_head等每 token-头动态 scale 类型以及 NVFP4 系列nvfp4、nvfp4_4over6可视为同一配置体系下的扩展面本文聚焦 FP8 主路径。CLI 与编程接口两种配置方式CLI 服务端使用缓存参数组中定义的--kv-cache-dtype见 arg_utils.py 参数组注册vllm serve model --kv-cache-dtype fp8离线/编程接口在构造LLM或AsyncLLM时传入同名关键字二者底层最终汇入同一个CacheConfigfrom vllm import LLM llm LLM( modelmeta-llama/Llama-3.1-8B-Instruct, kv_cache_dtypefp8, # 等价于 fp8_e4m3 )当kv_cache_dtypeauto时引擎启动阶段会通过resolve_kv_cache_dtype_stringarg_utils.py结合模型配置把 auto 解析为实际 dtype 后写入CacheConfig.cache_dtype。同时若模型 checkpoint 中自带 FP8 KV Cache 量化声明见下文 llm-compressor 校准且用户没有显式指定 dtypeattention.py 会自动把缓存 dtype 解析为fp8保证校准过就生效。两种 FP8 量化粒度Per-Tensor 与 Per-Attention-Head对同一个 FP8 缓存vLLM 支持两种量化/反量化 scale 粒度方案scale 形状说明Per-tensor 量化q/k/v_scale [1]Q、K、V 各自整体共享一个标量 scale开销最小精度相对粗Per-attention-head 量化q_scale [num_heads]k/v_scale [num_kv_heads]每个注意力头独立 scale精度更高需要保存/计算的头级 scale 张量更多其中 Per-attention-head 方案当前仅 FlashAttention 后端支持需要由llm-compressor提供的校准路径预先产出各头 scale无法在 vLLM 运行时凭空推断。从源码看Attention 层在初始化时会读取量化配置中的kv_cache_scheme当strategy attn_head时置use_per_head_quant_scales Trueattention.py后续将头级 scale 张量交由注意力后端使用在 FlashAttention 后端 中可以见到对layer._q_scale、layer._k_scale等做 shape 扩展expand(descale_shape)以匹配每头解量化的实现痕迹这正是每头 scale 落到 kernel 层的通路。Scale 校准的两种配置途径以及第三种隐式途径选哪个 dtype只是第一步更关键的问题是量化 scale 从哪里来。scale 若全为 1.0本质是截断式量化误差最大scale 越贴合真实激活分布精度越好。vLLM 提供以下方式方式一不做校准默认 scale 1.0仅通过kv_cache_dtypefp8开启量化所有 scale 一律取1.0llm LLM(modelmeta-llama/Llama-2-7b-chat-hf, kv_cache_dtypefp8)这种方式部署最简单、零额外依赖适合快速验证显存/吞吐收益但由于缺少与真实激活匹配的 scale精度损失通常最明显。引擎也会在 CacheConfig 校验 时打印对应提示。方式二用校准数据集估算 scale推荐对追求质量的生产场景vLLM 官方推荐使用llm-compressorvLLM 生态下的模型压缩库在一次性one-shot校准过程中用小规模有代表性的数据集统计出 Q/K/V 激活分布从而估计出最优 scale。它同时是开启 per-attention-head 高级量化的前提。具体操作见本文后面的端到端实操示例。第三种隐式途径直接消费 checkpoint 里固化好的 scalellm-compressor 校准完成后会把量化方案含kv_cache_scheme与 scale 系数随模型一起保存save_pretrained(..., save_compressedTrue)。之后在 vLLM 中加载该模型时如果用户没有显式传kv_cache_dtype即autoattention.py 会读取quant_config.kv_cache_scheme并自动切换到fp8缓存若 scheme 的 strategy 为attn_head则启用 per-head scale 通路同一文件 L288-L292。这意味着校准工作只需要做一次产出的模型即可在后续部署中自动复用 scale无需每次启动重复校准。对量化敏感层单独豁免--kv-cache-dtype-skip-layers并非所有注意力层都适合量化。例如sliding-window 注意力层常被用于长序列局部建模对 KV Cache 量化更敏感量化后质量下降更明显。kv_cache_dtype_skip_layers允许你把指定层留在模型原生 dtype而其余层继续使用所选量化 dtype实现量化与精度的折中。该参数在 CacheConfig 中定义为字符串列表接受两类输入注意力类型名如sliding_window表示匹配所有该类层层序号如0、1、23表示只豁免指定索引的层。命令行用法# 跳过所有 sliding-window 注意力层 vllm serve model \ --kv-cache-dtype fp8 \ --kv-cache-dtype-skip-layers sliding_window # 跳过第 0、1、23 层按层索引 vllm serve model \ --kv-cache-dtype fp8 \ --kv-cache-dtype-skip-layers 0 1 23编程接口用法from vllm import LLM llm LLM( modelmeta-llama/Llama-3.1-8B-Instruct, kv_cache_dtypefp8, kv_cache_dtype_skip_layers[sliding_window], )底层实现机制这个功能的判定逻辑位于 Attention 层初始化其执行流程为若该层存在 sliding-windowsliding_window is not None且 skip 列表包含sliding_window→ 命中否则用模块前缀名如model.layers.0.self_attn经 extract_layer_index 提取层索引若该索引以字符串形式出现在 skip 列表 → 命中命中后该层的kv_cache_dtype被改写为auto即回退到模型原生 dtype其他层不受影响。从 CacheConfig 的注释还能看到配套机制跳过层仍要参与 block-size 对齐必要时通过skip_page_size_padded把未量化层补齐到量化主层 page 大小保证混合缓存布局在页级可对齐。测试用例如tests/kernels/attention/test_attention.py对kv_cache_dtype相关矩阵的覆盖会验证 skip 列表与不同 dtype 组合的行为一致性。实操示例一无校准直接开启 FP8最简上手路径将示例模型改为你环境中实际可下载的模型即可仓库内示例使用 Llama-2-7B-chatfrom vllm import LLM, SamplingParams sampling_params SamplingParams(temperature0.7, top_p0.8) llm LLM( modelmeta-llama/Llama-2-7b-chat-hf, kv_cache_dtypefp8, ) prompt London is the capital of out llm.generate(prompt, sampling_params)[0].outputs[0].text print(out)运行后对比同模型auto与fp8两档fp8档的 KV Cache 占用显著下降可观测的峰值显存随之降低同时由于所有 scale 均为 1.0输出质量可能略降——这正是需要方式二校准的理由。实操示例二llm-compressor 数据集校准推荐高质量量化路径。以下脚本用HuggingFaceH4/ultrachat_200k的少量样本对 Llama 3.1 8B 做一次性校准将注意力激活Q与 KV Cache同时量化为 FP8并支持在tensor与attn_head两种 strategy 间切换。安装依赖pip install llmcompressorllm-compressor 还会依赖datasets、transformers等示例脚本 import 中使用如缺失请一并安装。校准脚本 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], tokenizeFalse) return tokenizer( text, paddingFalse, max_lengthMAX_SEQ_LEN, truncationTrue, add_special_tokensFalse, ) def build_recipe(strategy: str) - QuantizationModifier: fp8_args QuantizationArgs(num_bits8, typefloat, strategystrategy) return QuantizationModifier( config_groups{ attention: QuantizationScheme( targets[LlamaAttention], # Quantize queries: q_scale input_activationsfp8_args, ) }, kv_cache_schemefp8_args, # Quantize KV cache: k/v_scale ) # ----------------------------- # Main # ----------------------------- def main(): model AutoModelForCausalLM.from_pretrained(MODEL_ID, torch_dtypeauto) tokenizer AutoTokenizer.from_pretrained(MODEL_ID) ds load_dataset(DATASET_ID, splitf{DATASET_SPLIT}[:{NUM_CALIB_SAMPLES}]) ds ds.shuffle(seed42) ds ds.map( lambda ex: process_and_tokenize(ex, tokenizer), remove_columnsds.column_names, ) recipe build_recipe(STRATEGY) oneshot( modelmodel, datasetds, reciperecipe, max_seq_lengthMAX_SEQ_LEN, num_calibration_samplesNUM_CALIB_SAMPLES, ) save_dir f{MODEL_ID.rstrip(/).split(/)[-1]}-kvattn-fp8-{STRATEGY} model.save_pretrained(save_dir, save_compressedTrue) tokenizer.save_pretrained(save_dir) if __name__ __main__: main()脚本要点拆解QuantizationModifierconfig_groups把LlamaAttention的输入激活即 Q按fp8_args量化为 FP8产出 per-tensor 或 per-head 的q_scalekv_cache_schemefp8_args对 KV Cache 应用同一 FP8 量化参数产出k/v_scalestrategytensorvsattn_head前者每张 Q/K/V 张量一个 scale后者为每个头一个 scale量化粒度更细、精度更好但仅在 vLLM 的 FlashAttention 后端下被消费支持save_compressedTrue把量化配置含kv_cache_scheme与权重一起写入 checkpoint即上文第三种隐式途径所依赖的载体校准样本数与序列长度是质量-耗时权衡点NUM_CALIB_SAMPLES从 512 起步通常是不错的起点MAX_SEQ_LEN应贴近你的真实部署长度。在 vLLM 中加载校准产物校准并保存后即可在 vLLM 中直接加载该目录无需再次指定 FP8from vllm import LLM llm LLM(modelLlama-3.1-8B-Instruct-kvattn-fp8-tensor)引擎会自动识别 checkpoint 中固化的 FP8 KV Cache 方案并启用对应 attention.py 的 auto 解析路径。若要显式覆盖例如退回 bfloat16 验证基线则显式传kv_cache_dtypebfloat16此时用户显式取值优先于 checkpoint 声明。提示llm-compressor 的量化 KV Cache 示例会随其版本持续演进建议以其仓库的examples/quantization_kv_cache目录为最新参考本仓库仅说明 vLLM 侧的消费与配置语义。落到引擎内部缓存 dtype 是如何贯穿配置到 kernel 的把上文涉及的各条线索串起来一次 FP8 KV Cache 请求的完整链路如下参数层用户经LLM(...)或vllm serve --kv-cache-dtype ...传入kv_cache_dtype与kv_cache_dtype_skip_layers引擎在 create_engine_config 中将二者解析/透传给CacheConfig配置层CacheConfigcache.py持有cache_dtype在 validators 中依据 is_quantized_kv_cache 打印显存降低 可能精度损失的运行提示模型层每个 Attention 层初始化时按 skip 规则决定自身 dtype并结合 checkpoint 的kv_cache_scheme决定是否启用 per-head scaleattention.py后端层注意力后端按 dtype 选择 kernel 路径FP8 路径下 FlashAttention 类实现会把key_cache/value_cache以 FP8 语义读取并利用_q_scale/_k_scale做反量化对齐flash_attn.py。这一层抽象意味着用户在配置层看到的kv_cache_dtype最终是由平台能力CUDA/ROCm/HPU 各自支持的具体格式、Attention 后端能力是否支持 per-head scale与模型/校准产物共同决定的。选型与注意事项小结首次尝试用kv_cache_dtypefp8无校准先量级评估显存与吞吐收益同时用原模型做精度对照追求质量用 llm-compressor 数据集校准优先strategytensor需要更高精度时再切attn_head注意其 FlashAttention 后端限定与模型结构限定示例仅覆盖 Llama 类架构的LlamaAttention混合模型/长上下文若模型含 sliding-window 层用--kv-cache-dtype-skip-layers sliding_window或指定层索引做定点豁免平台前提fp8_e4m3需 CUDA 11.8 或 ROCmfp8_e5m2需 CUDA 11.8auto始终可用精度兜底量化缓存本质上是以精度换显存/吞吐务必在真实业务样本上回归评测量化不总是无损。更完整、随版本更新的校准示例请以 llm-compressor 官方示例为准vLLM 侧的完整参数语义可继续研读 CacheConfig、EngineArgs 与 Attention 层相关源码以及仓库内 test_attention.py、test_cache.py 等针对缓存 dtype 与 kernel 的测试用例。【免费下载链接】vllmA high-throughput and memory-efficient inference and serving engine for LLMs项目地址: https://gitcode.com/GitHub_Trending/vl/vllm创作声明：本文部分内容由AI辅助生成（AIGC），仅供参考
