APEX4: Efficient Pure W4A4 LLM Inference via Intra-SM Compute Rebalancing

2026-06-07

AI Summary

This work addresses the performance bottleneck of W4A4 quantization in practical deployment, where dequantization overhead on CUDA Cores often forces fallback to mixed precision. The study reveals for the first time that the feasibility of pure INT4 inference hinges on the hardware throughput ratio ρ between Tensor Cores and CUDA Cores, and proposes a ρ-aware pure INT4 inference method. By introducing a dedicated INT4 GEMM kernel, a granularity-adaptive mechanism, and group-wise dequantization optimization, the approach overcomes the dequantization bottleneck without modifying vLLM. On LLaMA-2-70B, it achieves perplexity close to FP16 and surpasses W4A8 Atom-g128 by 4.4% in zero-shot accuracy. End-to-end speedups reach up to 2.09× on L40S, RTX 3090, and A40 GPUs, and restore 1.4× performance on A100.

Abstract

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.

Problem

• W4A4 quantization
• dequantization overhead
• Tensor Cores
• CUDA Cores
• LLM inference

Innovation

• W4A4 quantization
• intra-SM compute rebalancing
• ρ-aware granularity adaptation
• pure INT4 inference
• Tensor Core/CUDA Core balance