Deployable W4A4 LLM Inference on Consumer Blackwell: Per-Token PE-QAT and Native INT4 mma.sync Kernels for SM121
Shresth Mishra
Abstract
Weight-activation quantization reduces memory and inference cost for LLMs, but practical W4A4 deployment remains difficult: activation outliers degrade accuracy, and low-bit inference paths designed for datacenter Blackwell do not carry over directly to consumer GB10/SM121 devices such as DGX Spark. This hardware distinction is fundamental rather than incremental: public SM100 Blackwell kernels rely on tcgen05 instructions and tensor memory (TMEM), whereas SM121 lacks these execution paths and instead requires a deployment strategy built around warp-level \texttt{mma.sync} and shared-memory tiling. This paper presents a deployment-oriented extension of Parameter-Efficient Quantization-Aware Training (PE-QAT) that reformulates the original per-channel W4A4 method into a per-token activation quantization pipeline compatible with native INT4 matrix multiply-accumulate on SM121. The method preserves the parameter-efficient training recipe of frozen base weights with LoRA adaptation while replacing non-deployable per-channel activation scaling with a single scale per token, enabling packed INT4 activations and weights to be consumed directly by an inline PTX \texttt{mma.sync.aligned.m16n8k64.row.col.s32.s4.s4.s32} kernel. A Hadamard variant further replaces learned smoothing and clipping with a fixed orthogonal rotation that diffuses activation outliers before per-token quantization and is calibration-free at deployment. The resulting runtime uses a two-kernel preprocessing plus INT4 tensor-core GEMM pipeline for multi-token inference and a dedicated INT4 GEMV path for single-token decode, applying floating-point scaling only after INT32 accumulation. On Llama-2-7B, the per-token PE-QAT variant achieves 67.11 average zero-shot accuracy across seven commonsense reasoning tasks, compared with 66.80 for FP16 and 66.69 for the original per-channel PE-QAT, while the Hadamard variant reaches 66.53 without deployment-time calibration. On DGX Spark GB10, both deployable variants reduce GPU memory from 15.9 GB to 7.6 GB, achieve 2.7--2.8$\times$ total speedup at 128-token prompt length and 2.2$\times$ at 512 tokens.
Chat is not available.
Successful Page Load