详解vLLM和SGLang awq dequantize kernel的魔法

文章预览

↑ 点击 蓝字  关注极市平台 作者丨BBuf 来源丨GiantPandaLLM 编辑丨极市平台 极市导读   本文详细解析了 vLLM 和 SGLang 中 AWQ int4 反量化 kernel 的实现原理和优化技巧。该 kernel 利用 IEEE 754 浮点数表示特性，通过 LOP3 和 PRMT 等 PTX 指令高效地将 int4 权重转换为 fp16 格式，避免了传统方法中的多次移位和类型转换，显著提升了反量化操作的性能。   >> 加入极市CV技术交流群，走在计算机视觉的最前沿 0x0. 前言 本片文章解析一下vLLM/SGLang中 awq int4的反量化kernel，这个kernel触发条件为当输入x的shape的tokens < 256时，这个时候会先把int4的awq权重使用 awq_dequantize 反量化回float16，然后调用PyTorch Matmul执行float16的乘法，代码位置见： https://github.com/vllm-project/vllm/blob/b82662d9523d9aa1386d8d1de410426781a1fa3b/vllm/model_executor/layers/quantization/awq.py#L162-L184 def apply(self,             l ………………………………