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

arXiv:2606.08761v1 Announce Type: cross
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 ($rho$) as the primary hardware indicator: the W4A4-g128 kernel yields $2.0$–$2.5times$ speedup on RTX~3090 ($rho=16$) yet degrades to $0.43$–$0.47times$ on A100 ($rho=64$) in compute-bond scenarios, establishing W4A4 viability as platform-dependent rather than universally infeasible. Guided by this finding, we build textbfAPEX4, which co-designs pure INT4 GEMM kernels with $rho$-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.66times$ end-to-end speedup on L40S ($rho=8$), and $1.78times$ on RTX~3090 ($rho=16$), $2.09times$ on A40 ($rho=16$), while recovering A100 ($rho=64$) to $1.20$–$1.40times$ via the mixed-granularity mode.