PaperlayerView exact revision on arXiv

Source-linked AI summary

Efficient INT8 Inference of Small NLP Models on Server CPUs with PyTorch Native Stack

Weiwen Xia, Yuxin Cui, E Cao

arXiv:2608.18182v1cs.CL

TL;DR

Small BERT-family models remain important for CPU-based production workloads, but native PyTorch quantization alone does not ensure fast enough inference. This work integrates and optimizes SmoothQuant in PyTorch for Intel Xeon CPUs, achieving up to 5.8× throughput speedup with negligible accuracy loss versus FP32.

  • Problem

    Native PyTorch quantization produces correct models but lacks backend-aware lowering and graph optimization needed for high-performance production CPU inference.

  • Method

    The paper integrates SmoothQuant into TorchAO and uses TorchInductor graph fusion, backend kernel selection, and Xeon-specific INT8 GEMM kernels to optimize inference.

  • Results

    Up to 5.8× end-to-end throughput speedup is achieved across BERT-family models and Xeon platforms, with negligible accuracy loss versus FP32 baselines.

  • Takeaways & Limitations

    The resulting workflow supports SmoothQuant-optimized INT8 inference for small NLP models directly through native PyTorch tooling with minimal deployment effort.

  • Takeaways & Limitations

    The underlying INT8 GEMM operator lacks native handling for blocked weight layouts, scales, and post-operations, requiring additional graph-level transformations.

Abstract

from arXiv · show

Small NLP models, especially BERT-family encoders, remain important in industrial workloads such as classification, ranking, and retrieval even in the era of large language models. On server CPUs, INT8 quantization offers an attractive latency-throughput-cost trade-off, but users increasingly expect such acceleration to be available directly in the native PyTorch stack. We integrate SmoothQuant into TorchAO and optimize the resulting inference path for Intel Xeon CPUs through graph-level fusion in TorchInductor and efficient INT8 GEMM kernel selection across oneDNN-, AVX512_VNNI-, and AMX-based implementations. Across BERT, DistilBERT, and XLM-RoBERTa benchmarks, the approach delivers up to 5.8x end-to-end throughput speedup with negligible---and in some cases no measurable---accuracy loss relative to the FP32 baseline. We also validated our work by detailed performance analysis with roofline models. The implementation has been upstreamed to PyTorch and TorchAO, enabling out-of-the-box deployment with native PyTorch tooling

1 Introduction

BERT-family small NLP models remain widely deployed on server CPUs, where INT8 quantization can reduce inference cost while preserving acceptable accuracy. This work integrates SmoothQuant into the native PyTorch stack and optimizes INT8 inference through TorchAO and TorchInductor for modern Xeon CPUs.

  • BERT-family small NLP models remain widely deployed for classification, ranking, and retrieval, while CPU inference offers favorable cost and broad infrastructure availability.
  • INT8 quantization targets transformer GEMM costs by reducing footprint and execution time while maintaining acceptable accuracy within the PyTorch ecosystem.
  • SmoothQuant is integrated into TorchAO’s frontend as an end-to-end workflow and improved to support subsequent graph optimizations.
  • TorchInductor graph fusion removes runtime overheads from INT8 GEMM weight-layout conversion and handles post-operations.
  • Benchmarks on BERT, DistilBERT, and XLM-RoBERTa show negligible accuracy loss and substantial throughput gains over FP32 baselines, supported by roofline analysis.
  • The resulting PyTorch-native path enables SmoothQuant-optimized INT8 deployment on server CPUs without leaving PyTorch or using third-party tools.

2 Related Work

Prior work spans integer-only quantization, PTQ, QAT, transformer-specific methods, and CPU inference software stacks. This work distinguishes itself by integrating SmoothQuant and PyTorch-native optimization components into one workflow for optimized INT8 inference.

  • Quantization methods: Integer-only quantization and subsequent PTQ and QAT surveys established the broader foundation for improving inference efficiency.These approaches frame quantization as a longstanding strategy for efficient inference.
  • Transformer quantization: Transformer quantization methods address activation outliers or sensitivity through mixed precision, weight-only quantization, activation-aware weighting, or shifting activation difficulty into weights.The cited methods include LLM.int8(), GPTQ, AWQ, and SmoothQuant.
  • CPU inference software: CPU inference stacks include Intel Neural Compressor, Hugging Face Optimum Intel, and ONNX Runtime, but they use differing deployment ecosystems and exported graphs.Intel Neural Compressor typically relies on PyTorch or third-party tools, Optimum Intel integrates with Hugging Face, and ONNX Runtime uses exported ONNX graphs.
  • Our approach: This work provides a fully native PyTorch workflow combining SmoothQuant, TorchAO, TorchInductor graph fusion, and ISA-aware CPU kernel selection.The integration enables seamless deployment of optimized INT8 inference within the PyTorch ecosystem.

3 Background and Challenges

This section explains how SmoothQuant makes transformer models more amenable to INT8 quantization while preserving linear-layer correctness, and why efficient CPU inference requires TorchInductor engineering and backend-aware kernel selection. It also distinguishes static and dynamic quantization and highlights support for both AVX512_VNNI-only and AMX-capable Xeon platforms.

  • Quantization Background: SmoothQuant redistributes scale factors between activations and weights to reduce activation outliers while preserving the linear layer mathematically.Activations are divided by per-channel smoothing factors, while corresponding weights are multiplied by them.
  • Quantization Background: Offline smoothing with a small calibration dataset makes the model substantially more amenable to standard INT8 quantization.The smoothing factors are derived before quantization using calibration data.
  • Quantization Background: Static quantization fixes scales and zero points during calibration with lower runtime overhead, whereas dynamic quantization computes them at runtime and may preserve accuracy more effectively at added execution cost.The choice depends on when activation quantization parameters are determined.
  • CPU Inference Challenges: TorchAO’s eager-mode quantization is correct but lacks backend-aware lowering and graph-level optimization, so production CPU inference requires compilation with TorchInductor.TorchInductor provides graph capture, optimization, and backend lowering; AOTI additionally removes Python overhead by compiling a shared binary.
  • CPU Inference Challenges: TorchInductor kernel selection is key to choosing the best backend implementation across different shapes, while supporting both AVX512_VNNI-only and AMX-capable Xeon platforms.The scaled INT8 GEMM operator is convenient for modeling and prototyping but has limited performance.

4 Design and Implementation

The implementation combines TorchAO SmoothQuant quantization with TorchInductor graph optimization and CPU-specific INT8 GEMM lowering. It fuses scaling and bias operations, selects between oneDNN and specialized AMX/AVX512_VNNI kernels, and adapts execution for Ice Lake’s unsigned-activation requirement.

  • TorchAO quantization: TorchAO prepares models for calibration, searches smoothing factors, stores quantized weights with scales, and replaces linear operations with INT8 × INT8 →INT32 GEMM followed by output scaling.After quantization, linear layers execute in INT8 while remaining operators stay in higher precision.
  • Kernel selection: TorchInductor’s CPP backend adds AMX and AVX512_VNNI INT8 microkernels, then generates, compiles, microbenchmarks, and selects the fastest candidate for future execution.This balances oneDNN generality against template-generated kernel shape specialization.
  • Graph fusion limitations: Runtime weight conversion, separate scale and post-operation handling, and oneDNN setup overhead limit the standalone torch._int_mm primitive.Efficient Xeon INT8 GEMM often requires blocked weight layouts tailored for VNNI and AMX instructions.
  • Graph fusion: TorchInductor fuses scaled INT8 GEMM subgraphs into oneDNN’s qlinear_pointwise operator, which accepts blocked weights, applies scales and bias-related post-operations, and caches metadata after warm-up.The fused operator can also be fused with other operators such as GeLU.
  • Ice Lake compatibility: On Ice Lake, activations are shifted from signed INT8 to unsigned INT8 for AVX512_VNNI, with zero-point compensation and weight conversion precomputed before inference.This enables u8 × s8 execution while preserving mathematical equivalence.

5 Performance Analysis and Optimization Targets

The analysis identifies linear-block kernels as the dominant BERT-large execution hotspot and uses roofline and Amdahl models to estimate INT8 acceleration headroom. INT8 GEMM is compute-bound on Ice Lake but memory-bound on Granite Rapids, yielding estimated end-to-end speedup ceilings of 3.15× and 7.06× versus FP32, respectively.

  • Baseline hotspots: About 87% of self CPU time on Granite Rapids and about 89% on Ice Lake is spent in linear-block kernels, versus 12% and 10% in attention.Linear-blocks approximate GEMM because profiler data cannot separate pure GEMM, and GEMM dominates these blocks.
  • Roofline analysis: FP32 GEMM is compute-bound on both platforms; BF16 GEMM is memory-bound on Granite Rapids; INT8 GEMM is compute-bound on Ice Lake but memory-bound on Granite Rapids.These classifications follow comparisons between GEMM arithmetic intensity and each platform’s data-type balance point.
  • Roofline analysis: 4.0× is the estimated GEMM speedup on Ice Lake, while Granite Rapids remains memory-bound for BF16 and INT8 despite INT8 weights fitting within LLC.The 442 MB INT8 model is smaller than the 504 MB LLC, but runtime-changing X and Y still require main-memory traffic.
  • Roofline analysis: After accounting for LLC-resident INT8 weights, INT8 GEMMs on Granite Rapids remain memory-bound, but effective memory bandwidth improves and the balance point decreases.The adjusted effective-bandwidth model combines main-memory traffic for activations and outputs with LLC traffic for weights.
  • End-to-end ceilings: Amdahl estimates INT8 end-to-end speedup ceilings of 3.15× versus FP32 on Ice Lake, 7.06× versus FP32 on Granite Rapids, and 1.53× versus BF16 on Granite Rapids.The estimates use GEMM fractions and speedups plus non-GEMM BF16 speedups of 1.20× on Ice Lake and 2.06× on Granite Rapids.
  • Caveats: The ceilings are idealized because quantization, data conversion, framework execution, post-operations, and uncertain LLC behavior add overhead or make linear-block latency approximate.The analysis explicitly notes that linear-block fractions are not pure GEMM fractions, so the end-to-end estimates are approximate.

6 Experiments and Evaluation

The experiments evaluate native SmoothQuant across BERT-family models, downstream accuracy tasks, and Intel Xeon platforms using the Section 5 deployment setup. The workflow delivers substantial INT8 speedups while maintaining accuracy with negligible loss.

  • Evaluation Setup: The evaluation reuses Section 5 hardware, sequence length, batch size, multi-instance strategy, AMP, and AOTI deployment settings, adding models, quantization methods, and downstream tasks.This extends baseline collection into full SmoothQuant benchmarking.
  • Models and Methods: BERT-large, DistilBERT, and XLM-RoBERTa are evaluated for performance and accuracy on SQuAD and MultiNLI using task-finetuned checkpoints.Both static and dynamic SmoothQuant are evaluated, with per-row quantization used for weights.
  • Performance Results: 4.2× to 5.8× throughput speedup is achieved over FP32 on Granite Rapids, while linear-block latency speedup exceeds 7.5×.On Ice Lake, throughput improves approximately 1.9× to 2.6× and linear-block latency decreases by more than 3×.
  • Ceiling Analysis: 2.58× throughput speedup is achieved by static quantization on Ice Lake, about 82% of the estimated ceiling, while linear-block latency speedup reaches 3.87×, about 97%.The realized Ice Lake speedup is lower than the ceiling, but the linear-block latency gap is relatively small because INT8 GEMM is compute-bound.
  • Performance Results: 5.82× throughput speedup over FP32 is achieved by static quantization on Granite Rapids, compared with 5.02× for dynamic quantization.The corresponding Granite Rapids linear-block latency speedups are about 8.7× for static quantization and 9.9× for dynamic quantization.
  • Accuracy Results: < 1% accuracy loss is observed with the proposed workflow, confirming substantial performance improvements with little loss in quality.Marginally higher quantized results are attributed to routine evaluation variance or mild overfitting rather than systematic quantization gains.

7 Conclusion

This work integrates and optimizes SmoothQuant-based INT8 inference for BERT-family models in native PyTorch on Intel Xeon CPUs. It extends TorchAO, expands TorchInductor graph fusion, and adds AVX512_VNNI- and AMX-based GEMM kernels for backend selection.

  • The workflow targets SmoothQuant-based INT8 inference for BERT-family NLP models in the native PyTorch stack on Intel Xeon CPUs.
  • TorchAO’s quantization workflow was extended to support SmoothQuant.
  • TorchInductor’s graph fusion was expanded to remove runtime overheads from layout conversion and separate post-operations.
  • AVX512_VNNI and AMX template-based GEMM kernels were added to participate in backend kernel selection.
Loading 2608.18182v1…