Jetson Orin AGX 上的 INT4 量化推理实践从 9 tok/s 到 19.6 tok/s项目地址https://github.com/luogantt/LLM-inference-engine这是一个从零实现的 C / CUDA 大模型推理引擎当前主要面向 DeepSeek-R1-Distill-Qwen-7B 的单 batch 推理。项目不依赖 vLLM、llama.cpp 或 PyTorch 推理框架核心 decode 路径由 CUDA 手写实现方便研究大模型推理中的真实瓶颈。这次重点是 Jetson Orin AGX 上的 INT4 量化推理优化。核心结论INT4 必须配合 int8 activation DP4A不能走 float 解包。如果只是把权重量化成 INT4然后在 kernel 里解包成 float 再做fmaf速度不会真正起来。INT4 的优势不只是权重体积小而是要把计算路径也切到整数点积上。测试环境Device: Jetson Orin AGX Model: DeepSeek-R1-Distill-Qwen-7B CUDA arch: sm_87 max_seq: 800 max_new_tokens: 512 batch: 1速度结果版本forward_msdecode speedFP16 baseline约 117.8 ms约 8.6 tok/sWeight-only INT8约 72-74 ms约 14.0 tok/sWeight-only INT4 float 解包约 109-111 ms约 9.3 tok/sINT4 weight INT8 activation DP4A约 52.7-53.3 ms约 19.6 tok/sINT4 DP4A 版本相对 FP1619.6 / 8.6 ≈ 2.28x相对 weight-only INT819.6 / 14.0 ≈ 1.40x相对旧的 INT4 float 解包路径19.6 / 9.3 ≈ 2.1x最初的问题一开始的 INT4 版本只做了 weight-only INT4FP16 weight - row-wise INT4 weight activation 保持 float kernel 内逐 byte 解包 int4 int4 - float float FMA理论上INT4 权重只有 FP16 的四分之一大小也只有 INT8 的一半大小应该能减少显存带宽压力。但实际结果并不好INT4 float 解包版本 ≈ 9.3 tok/s INT8 版本 ≈ 14.0 tok/s也就是说INT4 不但没有超过 INT8只是比 FP16 略快。问题原因原因在于这条路径省了带宽但增加了大量计算开销。旧 INT4 kernel 每个权重大致需要做这些事情load packed byte shift mask sign extend convert to float float fmaf这些操作发生在每个权重上。对于 7B 模型 decode 阶段的大量 GEMV 来说这个开销非常大。所以旧 INT4 路径的问题是权重确实变小了但每个权重都要运行时解包解包以后又回到 float FMA没有使用整数点积指令节省的显存带宽被解包和类型转换吃掉了这就是为什么 weight-only INT4 float 解包没有真正加速。正确方向新的 INT4 DP4A 路径改成weight: INT4 packed activation: 动态量化为 INT8 accumulate: int32 compute: __dp4a output: float核心变化是把计算路径从 float FMA 改成整数点积。__dp4a一次可以做 4 组 int8 乘加int8 x int8 - int32 accumulate由于 INT4 权重本身是 4-bit需要在寄存器里把 4 个 INT4 unpack 成 4 个 int8 lane然后和 int8 activation 做 DP4A。这条路径才真正符合 INT4 加速的逻辑更少的权重带宽 更紧凑的权重存储 整数点积计算 真实加速编译运行切换到 INT4 分支gitswitch jetson-orin-agx-int4gitpull编译 INT4 DP4A 版本make-fMakefile.cuda_lib clean-libmake-fMakefile.cuda_lib lib-int4Asm_87运行CUDA_VISIBLE_DEVICES0python python_infer.py\--model/data/project/deepseek-r1-7b\--lib./build/libllm_cuda.so\--prompt你好 deepseek 介绍一下黑格尔的思想\--max-new-tokens512\--max-seq800如果需要对比旧的 float 解包 INT4 路径make-fMakefile.cuda_lib clean-libmake-fMakefile.cuda_lib lib-int4-floatAsm_87当前实现方式当前项目仍然使用原始 HuggingFace safetensors 模型目录。INT4 不是读取 GPTQ/AWQ 格式而是在加载权重时做 row-wise INT4 量化FP16/BF16 safetensors - row-wise INT4 packed weight每一行权重保存一个 scaleweight_float ≈ int4_value * weight_scaleactivation 在每次 linear 前动态量化为 INT8activation_float ≈ int8_value * activation_scale最终输出output_float ≈ int32_accumulate * weight_scale * activation_scale为什么这个结果有意义Jetson Orin AGX 是边缘端设备显存带宽和功耗都比数据中心 GPU 更敏感。单 batch decode 又是典型的 memory-bound 场景大量计算以 GEMV 为主。因此INT4 的价值不只是省显存更重要的是让 decode 路径更适合边缘端硬件模型权重更小显存带宽压力更低DP4A 可以利用整数点积单 token decode 延迟明显下降在这个实验中INT4 DP4A 让 DeepSeek-R1-Distill-Qwen-7B 在 Jetson Orin AGX 上达到约 19.6 tok/s已经明显超过 FP16 和 weight-only INT8。后续优化方向当前版本还有继续优化空间把 activation quantize 融合进 linear kernel减少额外 kernel launch针对HIDDEN3584和INTERMEDIATE18944做专用 GEMV kernel减少 QKV、MLP 中重复的量化开销对输出层lm_head做单独优化尝试 group-wise scale提高精度和速度平衡结合 CUDA Graph 降低 decode step 的调度开销总结这次实验最重要的经验是INT4 不是只把权重压到 4bit 就会快。 INT4 必须让计算路径也进入整数点积。 INT4 必须配合 int8 activation DP4A不能走 float 解包。项目还在继续优化中欢迎关注和交流https://github.com/luogantt/LLM-inference-engine