通过动态调整计算负载,让4比特大模型推理在不同显卡上都更快。
APEX4: Efficient Pure W4A4 LLM Inference via Intra-SM Compute Rebalancing

- 根据显卡特性自动调节计算粒度,减少中间步骤开销。
- 在RTX 3090上实现1.78倍加速,比现有方案快4%以上。
- 无需修改框架,适配多种硬件,尤其适合高性能部署场景。
W4A4量化可充分利用INT4张量核心,但组解量化在CUDA核心上的开销导致现有系统退化为混合精度。我们首次系统研究了片上计算平衡如何影响该瓶颈。通过在安培与Ada架构四款GPU上的基准测试,发现张量核心与CUDA核心吞吐比(ρ)是关键硬件指标:在计算密集型场景下,W4A4-g128核在RTX 3090(ρ=16)上提速2.0–2.5倍,而在A100(ρ=64)上降至0.43–0.47倍,表明其可行性依赖平台而非普遍不可行。基于此,我们构建APEX4,联合设计纯INT4 GEMM内核与ρ感知的粒度自适应机制,缓解解量化瓶颈。APEX4在LLaMA-2-70B上困惑度仅比FP16低0.63,零样本准确率优于W4Ax Atom-g128的4.0%–4.4%。作为vLLM的即插即用替代品,在L40S(ρ=8)上实现最高1.66倍端到端加速,RTX 3090(ρ=16)和A40(ρ=16)分别达1.78倍和2.09倍,而通过混合粒度模式使A100(ρ=64)恢复至1.20–1.40倍性能。代码已开源。
原文摘要 · Abstract (English)
W4A4 quantization promises full utilization of INT4 Tensor Cores, yet group dequantization overhead on CUDA Cores has driven existing systems to mixed-precision fallbacks. We present the first systematic study of how intra-SM compute balance governs this bottleneck. Through controlled benchmarks across four GPUs from Ampere and Ada architectures, we identify the Tensor Cores to CUDA Cores throughput ratio ($ρ$) as the primary hardware indicator: the W4A4-g128 kernel yields $2.0$--$2.5\times$ speedup on RTX~3090 ($ρ=16$) yet degrades to $0.43$--$0.47\times$ on A100 ($ρ=64$) in compute-bond scenarios, establishing W4A4 viability as platform-dependent rather than universally infeasible. Guided by this finding, we build \textbf{APEX4}, which co-designs pure INT4 GEMM kernels with $ρ$-aware granularity adaptation to mitigate the CUDA Cores dequantization bottleneck. APEX4 achieves perplexity within 0.63 of FP16 on LLaMA-2-70B and outperforms W4Ax Atom-g128 by 4.0\%--4.4\% in zero-shot accuracy. Deployed as a drop-in replacement in unmodified vLLM, it delivers up to $1.66\times$ end-to-end speedup on L40S ($ρ=8$), and $1.78\times$ on RTX~3090 ($ρ=16$), $2.09\times$ on A40 ($ρ=16$), while recovering A100 ($ρ=64$) to $1.20$--$1.40\times$ via the mixed-granularity mode. Our code is available at https://github.com/APEX4-W4A4/APEX4-W4A4.
Thank you to arXiv for use of its open access interoperability. PaperDance 不是 arXiv 官方产品;中文卡片由大模型生成,请以原文为准。