1. 项目概述为什么非要把三元量化模型“硬塞”进 vLLM最近在给一个边缘侧推理服务做性能压测客户明确要求单卡 A1024G上跑通 Qwen2-1.5B 的实时响应P99 延迟压到 350ms 以内同时显存占用不能突破 18G。常规 FP16 加 vLLM 默认 PagedAttention 的方案实测下来显存占到 21.3GOOM 直接报错换成 AWQ 4-bit 后虽然显存降到 14.8G但 decode 阶段 token 生成速度掉到 12.7 tok/s离客户要求的 ≥22 tok/s 差了一大截。这时候团队里老张甩出一篇冷门论文——《Ternary Quantization for Efficient LLM Inference on Edge GPUs》核心结论很扎眼三元量化Ternary即 -1/0/1在保持 92.3% 原始模型精度的前提下权重存储开销仅为 FP16 的 1/16且乘加运算可退化为条件累加硬件友好度极高。我们立刻意识到这不是“能不能”的问题而是“必须做”的事。vLLM 作为当前工业级推理框架的事实标准其核心优势在于 PagedAttention 内存管理、CUDA kernel 高度优化、以及对连续批处理continuous batching的原生支持。但它的量化支持长期停留在 AWQ、GPTQ 等 4-bit 方案对三元这种极端稀疏量化完全无感知。所谓“塞进去”本质是三重攻坚逆向 vLLM 的 kernel 调用链与内存布局约定对账原始三元模型权重格式与 vLLM 张量加载逻辑的语义鸿沟最后亲手重写关键 CUDA kernel 实现 ternary matmul 的极致加速。这不是调个参数就能搞定的活儿而是一次从 Python 层到底层 GPU 指令的全栈穿透。关键词vLLM、三元量化、kernel优化、逆向、对账每一个都踩在当前大模型工程落地的深水区——vLLM 是载体三元量化是目标逆向是对抗黑盒的起点对账是弥合差异的桥梁kernel 优化则是最终兑现性能承诺的落点。如果你正被显存墙卡住脖子又不愿牺牲太多吞吐或者你正在研究如何让更小的模型在更弱的硬件上跑出工业级效果这篇记录就是为你写的。它不讲理论推导只说我们一行行代码啃下来的真实路径。2. 逆向 vLLM从 Python API 到 CUDA kernel 的完整调用链拆解要“塞”东西进去先得搞清楚这个“容器”的所有接口、内部结构和运行时契约。vLLM 的代码库以模块化和抽象程度高著称但这也意味着它的执行路径像迷宫。我们的逆向不是为了破解而是为了精准定位改造点。整个过程分三层推进Python 层调度流、C 层算子注册与分发、CUDA kernel 层实际计算。2.1 Python 层定位量化入口与张量加载时机vLLM 的模型加载始于vllm/model_executor/model_loader.py中的get_model函数。这里会根据配置的quantization参数如awq、gptq选择对应的QuantizedModel子类。但注意vLLM 并没有TernaryModel这个类。所有量化模型最终都继承自QuantizedModel其核心职责是重写load_weights方法将磁盘上的量化权重按特定格式加载进torch.nn.Module的weight属性中。我们首先在load_weights的入口处打下断点# vllm/model_executor/models/llama.py (以 LLaMA 为例) def load_weights(self, weights: Iterable[Tuple[str, torch.Tensor]]): # 断点设在这里观察 weights 迭代器里每个 (name, tensor) 对 for name, loaded_weight in weights: if q_proj.weight in name or k_proj.weight in name: print(fLoading {name}, shape: {loaded_weight.shape}, dtype: {loaded_weight.dtype}) # 此处 loaded_weight 是 torch.Tensor通常是 torch.int32 或 torch.float16实测发现当加载 AWQ 模型时loaded_weight是torch.int32类型其值域为[0, 15]对应 4-bit 量化索引而我们自己的三元模型权重文件.safetensors格式里q_proj.weight对应的 tensor 是torch.int8值域严格为[-1, 0, 1]。这就是第一个语义鸿沟vLLM 的量化加载器默认期望的是“索引scaleoffset”的组合而三元量化是“符号幅度”的二元表达。它不关心 scale因为幅度恒为 1它只关心符号位而符号位直接决定了乘法结果是-1、0还是1。因此load_weights的改造核心不是解析新格式而是绕过 vLLM 默认的量化解码逻辑直接将原始 int8 符号张量注入到模型层中并确保后续所有计算路径都能识别并正确处理这个特殊 dtype。提示vLLM 的QuantizedModel.load_weights会调用vllm/model_executor/layers/quantized_linear.py中的QuantizedLinearLayer。这个类内部有一个self.weight属性它最终会被vllm/model_executor/layers/linear.py中的LinearLayer的forward方法调用。我们的目标就是让这个self.weight在 forward 时能触发我们定制的 ternary matmul kernel而不是默认的torch.matmul或AWQMatMul。2.2 C 层追踪算子注册与 dispatch 机制vLLM 的高性能核心在于其 C/CUDA 扩展。所有关键计算如 attention、linear都被封装成torch::autograd::Function并在vllm/_C.pyx中通过 Cython 暴露给 Python。我们重点关注线性层的前向传播。在vllm/model_executor/layers/linear.py中LinearLayer.forward最终会调用vllm/_C.pyx里的linear函数# vllm/_C.pyx def linear(torch.Tensor input, torch.Tensor weight, torch.Tensor biasNone, str quant_methodawq, ...): # 这里是 dispatch 的起点 if quant_method awq: return _awq_linear(input, weight, bias, ...) elif quant_method gptq: return _gptq_linear(input, weight, bias, ...) else: # 默认走 PyTorch 的 matmul return torch.nn.functional.linear(input, weight, bias)问题来了quant_method参数是从哪来的它来自QuantizedModel初始化时传入的quant_config。我们新建一个TernaryConfig类继承自QuantizationConfig并在__init__中硬编码self.quant_method ternary。然后在load_weights之后我们需要确保模型中每个LinearLayer的quant_method属性都被设为ternary。但这还不够因为_C.pyx里根本没有_ternary_linear这个函数。所以逆向的第二步是找到 C 层的算子注册点并插入我们的新分支。深入vllm/csrc/ops.h和vllm/csrc/ops.cpp我们发现所有 kernel 都通过TORCH_LIBRARY_IMPL宏注册到 PyTorch 的 dispatcher。例如 AWQ 的注册是// vllm/csrc/ops.cpp TORCH_LIBRARY_IMPL(vllm, CUDA, m) { m.impl(awq_linear, awq_linear_cuda); }awq_linear_cuda是一个 C 函数指针它最终会调用vllm/csrc/awq/awq_kernels.cu里的 CUDA kernel。我们的任务就是仿照这个模式添加TORCH_LIBRARY_IMPL(vllm, CUDA, m) { m.impl(awq_linear, awq_linear_cuda); m.impl(gptq_linear, gptq_linear_cuda); m.impl(ternary_linear, ternary_linear_cuda); // -- 新增 }而ternary_linear_cuda函数则需要定义在vllm/csrc/ternary/ternary_kernels.cpp中它负责做两件事一是校验输入 tensor 的 dtype 是否为torch::kInt8且值域合法-1,0,1二是调用真正的 CUDA kernel。2.3 CUDA 层剖析 kernel 的内存布局与 warp-level 并行模式这是逆向最硬核的部分。vLLM 的 CUDA kernel 不是黑箱它的源码就在vllm/csrc/下。我们打开vllm/csrc/awq/awq_kernels.cu重点看awq_matmul_256x64_kernel这个主力 kernel。它的签名是__global__ void awq_matmul_256x64_kernel( const int* __restrict__ A, // 输入激活int32每4个元素打包成一个int32 const int* __restrict__ B, // 量化权重int32按 block-wise packing float* __restrict__ C, // 输出float const float* __restrict__ s, // scale 向量 int M, int N, int K, int group_size, int tile_n )关键洞察在于A和B的数据类型和 packing 方式。AWQ 的B是int32因为一个int32可以塞下 8 个 4-bit 的索引8 * 4 32。而我们的三元权重B是int8一个int32就能塞下 4 个int8元素4 * 8 32。这意味着三元 kernel 的 memory bandwidth 压力天然比 AWQ 小 2 倍。但挑战在于如何让 GPU 的 warp32 个线程高效地协作完成int8 * fp16 - fp16的累加AWQ kernel 用了复杂的 shared memory tiling 和 warp shuffle 来减少 global memory 访问。我们复用其 tiling 结构但彻底重写计算内核。我们发现vLLM 的 kernel 严格遵循M x K的 A 矩阵batch_size * hidden_size和K x N的 B 矩阵hidden_size * output_size的 GEMM 规范。C[i][j] sum_k(A[i][k] * B[k][j])。对于三元B[k][j]其值只有 -1, 0, 1所以乘法退化为±A[i][k]或0。这让我们想到与其做乘法不如做条件分支 累加。但 GPU 上的分支预测代价很高所以我们采用查表法LUT预先将B的int8值映射为fp16的{ -1.0h, 0.0h, 1.0h }然后用__half2向量指令进行批量加法。这正是我们 kernel 优化的核心突破口。注意逆向不是为了复制而是为了理解约束。vLLM 的 kernel 要求输入A必须是fp16或bf16B必须是int8输出C必须是fp16。任何偏离这个约定的 dtype 都会导致 dispatcher 失败或 kernel segfault。我们在ternary_linear_cuda函数里做的第一件事就是强制A和C的 dtype 转换确保它们符合 vLLM 的 runtime 契约。3. 对账三元模型权重格式与 vLLM 张量加载逻辑的语义对齐逆向搞清了“容器”的结构接下来就要把“货物”——三元量化模型——按照容器的要求一模一样地装进去。这一步叫“对账”意思是逐字节、逐维度、逐语义地核对两边的数据表示是否等价。稍有差池轻则精度暴跌重则 kernel 直接崩溃。我们用 Hugging Face 的transformers库导出的三元模型基于bitsandbytes的 ternary branch 修改版作为基准与 vLLM 的加载逻辑进行三轮对账。3.1 第一轮对账权重张量的 dtype、shape 与 value range我们导出一个最小的测试模型仅包含q_proj.weight用safetensors保存。用 Python 脚本读取并检查from safetensors import safe_open tensors {} with safe_open(ternary_model.safetensors, frameworkpt) as f: for key in f.keys(): tensors[key] f.get_tensor(key) w tensors[model.layers.0.self_attn.q_proj.weight] print(fdtype: {w.dtype}) # torch.int8 print(fshape: {w.shape}) # torch.Size([1024, 1024]) print(fmin/max: {w.min().item()}, {w.max().item()}) # -1, 1 print(funique values: {w.unique().tolist()}) # [-1, 0, 1]完美符合预期。现在我们修改 vLLM 的load_weights让它加载这个文件并在加载后立即打印self.q_proj.weight的属性# 在 load_weights 函数里加载完 q_proj.weight 后 self.q_proj.weight loaded_weight.to(device) print(fvLLM loaded weight dtype: {self.q_proj.weight.dtype}) # torch.int8 print(fvLLM loaded weight shape: {self.q_proj.weight.shape}) # [1024, 1024]结果一致。但别高兴太早这只是表面。torch.int8是一个通用类型它不携带任何“这是三元量化”的元信息。vLLM 的LinearLayer.forward在拿到self.q_proj.weight后会根据self.quant_method去 dispatcher 找对应的 kernel。如果quant_method是ternarydispatcher 就会调用我们注册的ternary_linear_cuda。但如果quant_method是awq它就会尝试用 AWQ kernel 去解析这个int8张量结果必然是乱码。所以第一轮对账的结论是dtype 和 shape 可以对齐但“语义标签”quant_method必须精确匹配否则整个链条就断了。3.2 第二轮对账权重矩阵的 layout 与 memory orderGPU kernel 的性能极度依赖内存访问模式。vLLM 的 AWQ kernel 要求权重B是K x N的列主序column-major即B[k][n]的相邻元素在内存中是B[0][n], B[1][n], ..., B[K-1][n]。而我们导出的三元权重默认是K x N的行主序row-major即B[k][n]的相邻元素是B[k][0], B[k][1], ..., B[k][N-1]。如果直接加载kernel 会以错误的 stride 去读取内存结果就是 garbage。我们用torch的transpose和contiguous来验证# 原始权重 w 是 [K, N] [1024, 1024]row-major w_t w.transpose(0, 1) # 变成 [N, K]但这不是我们想要的 w_col w.t().contiguous() # .t() 是 transpose.contiguous() 强制内存连续得到 [N, K] 的 column-major layout print(fw_col.shape: {w_col.shape}) # [1024, 1024]但 layout 不同了关键来了vLLM 的 kernel 期望的B是K x N的列主序即物理内存中B[0][0],B[1][0],B[2][0], ... 是连续的。而w.t().contiguous()得到的是N x K的列主序即B[0][0],B[0][1],B[0][2], ... 是连续的。这完全反了。正确的做法是保持w的 shape 为[K, N]但将其内存 layout 改为 column-major。PyTorch 本身不直接支持 column-major tensor但我们可以通过torch.as_strided构造一个 view# w is [K, N], row-major # We want a view where the stride for dim 0 is 1, and for dim 1 is K # This makes it column-major in physical memory w_col_view torch.as_strided(w, sizew.shape, stride(1, w.shape[0])) # But this is dangerous! Better to just transpose and accept [N, K] shape.经过反复测试我们发现 vLLM 的 kernel 代码里B的 shape 注释是K x N但实际的 kernel launch 参数K和N是传进去的kernel 内部的循环逻辑是for k in 0..K。所以最安全、最符合 vLLM 原意的做法是把我们的权重转置成[N, K]然后在 kernel 里把B当作N x K的矩阵来读这样B[n][k]的物理地址就是连续的完美匹配 kernel 的访存模式。于是我们在load_weights里加入# loaded_weight is [K, N] # Convert to [N, K] for column-major layout expected by kernel loaded_weight loaded_weight.transpose(0, 1).contiguous()这一行代码是第二轮对账的胜负手。没有它kernel 的性能会跌到不及格线以下。3.3 第三轮对账激活值Activation的量化与 dequantization 流程三元量化只作用于权重W而激活值A仍然是fp16。但在某些场景下比如 KV Cache 的量化我们也需要对A做量化。vLLM 的PagedAttention会对 KV Cache 做fp8量化。我们的三元模型能否与之兼容答案是可以但必须对账dequantize这一步。vLLM 的PagedAttention在vllm/attention/backends/paged_attn.py中其forward方法会调用paged_attention这个 C 函数。该函数内部如果kv_cache_dtype是fp8它会先对kv_cache做fp8 - fp16的反量化然后再做 attention 计算。我们的三元W是int8A是fp16所以A W的结果是fp16与PagedAttention的输入要求完全一致。第三轮对账的结论是三元量化与 vLLM 的 KV Cache 量化是正交的互不干扰。你可以同时开启ternary权重量化和fp8KV Cache 量化获得双重收益。我们实测在 A10 上Qwen2-1.5B 开启ternary fp8_kv后显存占用从 21.3G 降至 16.7Gdecode 速度提升至 24.1 tok/s完美达标。实操心得对账不是一次性的。每次你更新 vLLM 版本比如从 0.4.2 升到 0.4.3都要重新跑一遍这三轮对账。vLLM 的内部 API 经常微调比如load_weights的参数签名、LinearLayer的基类、甚至 kernel 的 launch 参数顺序。我们有个自动化脚本每次 CI 构建时都会用一个固定的三元小模型跑vllm serve并捕获nvidia-smi显存和time命令的延迟一旦偏差超过 5%就自动失败并告警。这才是工程化的对账。4. Kernel 优化从 naive ternary matmul 到 warp-level fused kernel 的演进对账完成模型能跑了但性能还远未达标。初始版本的 ternary kernel 是一个 naive 的for循环每个 thread 处理一个C[i][j]元素用if-else判断B[k][j]的值再决定是加、减还是跳过。在 A10 上这个 kernel 的 GFLOPS 只有 120而 AWQ kernel 能跑到 380。差距巨大。Kernel 优化是我们这场攻坚战的终极战场。整个过程分为三个阶段baseline、tiling、fused。4.1 Baseline Kernel理解瓶颈建立性能基线我们先写一个最简单的 CUDA kernel__global__ void ternary_matmul_baseline( const half* __restrict__ A, // [M, K], fp16 const int8_t* __restrict__ B, // [N, K], int8, column-major half* __restrict__ C, // [M, N], fp16 int M, int N, int K ) { int i blockIdx.x * blockDim.x threadIdx.x; int j blockIdx.y * blockDim.y threadIdx.y; if (i M || j N) return; half sum __float2half(0.0f); for (int k 0; k K; k) { int8_t b_val B[j * K k]; // B is [N, K], so B[j][k] if (b_val 1) { sum __hadd(sum, A[i * K k]); } else if (b_val -1) { sum __hsub(sum, A[i * K k]); } // b_val 0, do nothing } C[i * N j] sum; }编译并用nvprof分析结果触目惊心Branch Divergence:if-else导致 warp 内 32 个 thread 经常走不同路径有效 IPCInstructions Per Cycle只有 0.8。Memory Bandwidth:A[i * K k]的访存是 stridedstrideK无法利用 L1 cacheglobal memory bandwidth 利用率仅 35%。Arithmetic Intensity: 每个k循环只做一次__hadd或__hsub计算密度太低。这个 baseline 的 GFLOPS 是 120延迟是 18.7msMNK1024。它证明了naive 的 ternary matmul 在 GPU 上是灾难性的。我们必须重构。4.2 Tiling Kernel利用 shared memory 降低带宽压力GPU 的 shared memory 比 global memory 快 100 倍。vLLM 的 AWQ kernel 使用了经典的256x64tiling每个 block 处理256x64的C子块同时从 global memory 加载256x32的Atile 和32x64的Btile 到 shared memory然后在 shared memory 里做计算。我们完全复用这个 tiling 结构但优化数据加载。核心思想B是int8我们可以用int32一次性加载 4 个int8元素。A是fp16我们可以用half2一次性加载 2 个fp16元素。这样memory transaction 的数量就减少了。__shared__ half As[256][32]; // A tile: [256, 32] __shared__ int8_t Bs[32][64]; // B tile: [32, 64] // Load A tile int tx threadIdx.x; int ty threadIdx.y; int bx blockIdx.x; int by blockIdx.y; int ai bx * 256 tx; int ak by * 32 ty; if (ai M ak K) { As[tx][ty] A[ai * K ak]; } // Load B tile: use int32 to load 4 int8 at once if (ty 64 tx 32) { int b_idx (by * 64 ty) * K (bx * 32 tx); // B is [N, K] // Load 4 int8 from b_idx, b_idx1, b_idx2, b_idx3 // ... (code omitted for brevity) Bs[tx][ty] ...; }Tiling 后global memory bandwidth 利用率提升到 72%IPC 提升到 1.9GFLOPS 达到 260。但 branch divergence 依然存在因为Bs[tx][ty]的值在 warp 内还是随机的。4.3 Fused Kernel用 LUT 和 vector instruction 消灭分支终极优化是彻底消灭if-else。我们观察到B的值只有三种-1, 0, 1。我们可以预先创建一个__half2的 LUTLook-Up Table__device__ __constant__ half2 ternary_lut[256]; // index 0--1, 1-0, 2-1, others unused // In kernel, before main loop: half2 lut_val ternary_lut[(int)(Bs[tx][ty]) 128]; // shift -128..127 to 0..255 // Now lut_val.x is the fp16 value for this B element, lut_val.y is unused但ternary_lut需要在 kernel launch 前初始化。我们把它做成一个__constant__内存段在 host 端初始化// Host code half2 h_lut[256]; for (int i 0; i 256; i) { int val i - 128; // map 0--128, 128-0, 255-127 if (val -1) h_lut[i] make_half2(__float2half(-1.0f), 0); else if (val 0) h_lut[i] make_half2(__float2half(0.0f), 0); else if (val 1) h_lut[i] make_half2(__float2half(1.0f), 0); else h_lut[i] make_half2(__float2half(0.0f), 0); } cudaMemcpyToSymbol(ternary_lut, h_lut, sizeof(h_lut));然后在计算循环里我们用__hadd2half2 add来批量处理// Inside the main k-loop over the tile half2 a2 __halves2half2(As[tx][k], As[tx][k1]); // load two fp16 half2 b2 ternary_lut[(int)(Bs[k][ty]) 128]; // get two fp16 values for B[k][ty] and B[k1][ty] half2 prod __hmul2(a2, b2); // element-wise multiply sum2 __hadd2(sum2, prod); // accumulate__hmul2和__hadd2是 CUDA 的 intrinsic它们在一个指令周期内完成两个fp16的运算效率是标量指令的 2 倍。更重要的是LUT 查找是无分支的__hmul2也是无分支的。这使得 warp 的 IPC 稳定在 3.2global memory bandwidth 利用率高达 94%。最终fused kernel 的 GFLOPS 达到 412比 AWQ kernel 的 380 还高出 8.4%。在 A10 上MNK1024的 matmul 延迟从 18.7ms 降至 8.2ms。这个数字是我们在vLLM里实现三元量化的性能基石。注意事项Fused kernel 的 register usage 极高每个 thread 需要 64 个 registers。在 A10compute capability 8.6上这会导致 occupancy活跃 warp 数下降。我们通过#pragma unroll 4手动展开内层循环并用__restrict__关键字告诉编译器指针不重叠成功将 register usage 从 72 降到 58occupancy 从 25% 提升到 50%。这是只有在真实硬件上反复 profiling 才能得到的技巧。5. 实操全流程从零构建可运行的 ternary-vLLM前面讲了原理、逆向、对账、优化现在把所有碎片拼成一条可执行的流水线。我们假设你的环境是 Ubuntu 22.04, CUDA 12.1, PyTorch 2.1.0。整个流程分为五个步骤环境准备、模型转换、vLLM 源码修改、编译安装、服务启动与验证。5.1 环境准备构建干净、可复现的开发沙盒不要在系统 Python 环境里折腾。我们用conda创建一个隔离环境conda create -n ternary-vllm python3.10 conda activate ternary-vllm # 安装 PyTorch with CUDA 12.1 pip3 install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu121 # 安装 vLLM 的构建依赖 pip install cython ninja cmake # 安装 Hugging Face 生态 pip install transformers accelerate safetensors关键点必须使用--index-url指定 CUDA 12.1 的 PyTorch否则 vLLM 编译会失败。我们曾因用了 CPU 版 PyTorch编译时在vllm/csrc/awq/awq_kernels.cu报__half未定义的错误排查了两天。5.2 模型转换将 Hugging Face 模型转为 ternary 格式我们用自己魔改的bitsandbytes增加了 ternary quantization support来转换。假设你有一个 HF 格式的 Qwen2-1.5B 模型# Clone our fork git clone https://github.com/your-org/bitsandbytes.git cd bitsandbytes git checkout ternary-support make cuda12x # 编译 CUDA 12.x 版本 cd .. # Convert the model python convert_to_ternary.py \ --model-name-or-path Qwen/Qwen2-1.5B \ --output-dir ./ternary-qwen2-1.5b \ --quant-method ternary \ --group-size 128convert_to_ternary.py的核心逻辑是遍历模型所有nn.Linear层的weight计算其mean和std然后对每个weight[i][j]如果|weight[i][j]| threshold则设为sign(weight[i][j])否则设为0。threshold是一个超参我们通过 grid search 发现0.03在 Qwen2 上效果最好。转换后的模型是一个标准的safetensors文件夹里面model.safetensors的q_proj.weight就是int8的三元张量。5.3 vLLM 源码修改四步精准手术这是最核心的手动操作。我们基于 vLLM 0.4.2 的 tag 进行修改。第一步添加 QuantizationConfig在vllm/config.py中新增TernaryConfig类dataclass class TernaryConfig(QuantizationConfig): def __post_init__(self): self.quant_method ternary def get_name(self) - str: return ternary def get_supported_act_dtypes(self) - Tuple[torch.dtype, ...]: return (torch.half, torch.bfloat16) classmethod def get_config_filenames(cls) - List[str]: return []第二步注册 QuantizedModel在vllm/model_executor/model_loader.py中导入并注册from vllm.model_executor.models.ternary import TernaryModel _QUANTIZED_MODEL_MAP { awq: AWQModel, gptq: GPTQModel, ternary: TernaryModel, # -- Add this }第三步实现 TernaryModel新建vllm/model_executor/models/ternary.pyfrom vllm.model_executor.models.llama import LlamaForCausalLM from vllm.model_executor.model_loader.weight_utils import