PaperlayerView exact revision on arXiv

Source-linked AI summary

A Method for Layer Bit-Width Allocation in LLM Quantization via Performance Maximization Under a Quality-Degradation Constraint

Artem Safronov

arXiv:2608.28003v1cs.LGcs.AI

TL;DR

The paper addresses how to allocate quantization precision across Gemma-3-1B layers to improve inference speed without exceeding a quality-loss budget. It uses SA-PTQ sensitivity estimates with TensorRT-LLM to test blockwise W8A8 configurations, finding that FFN and lm_head quantization can accelerate inference while short-context Attention quantization can reduce throughput. The best low-degradation configuration reduces latency by 11.0%, while a higher-degradation configuration reaches 19.1%.

  • Problem

    Uniform quantization ignores differing layer sensitivities, while mixed-precision transitions introduce overhead that requires direct measurement under a controlled quality-loss constraint.

  • Method

    The paper uses SA-PTQ sensitivity estimates and TensorRT-LLM activation pass-through to allocate INT8 precision blockwise across FFN, Attention, and lm_head subsystems.

  • Results

    11.0% latency reduction: FFN 5+5 + lm_head with 98.90% Top-1 agreement and +0.85% perplexity degradation; FFN all26 + lm_head reaches up to 19.1% speedup with noticeable quality loss.

  • Takeaways & Limitations

    FFN and lm_head are practical W8A8 acceleration targets, whereas short-context Attention needs more effective fused INT8 kernels or partial quantization to realize additional gains.

  • Takeaways & Limitations

    The Attention analysis covers short contexts; longer contexts and accumulated KV-cache may substantially change its behavior, and projection-specific degradation was not analyzed by SQNR.

Abstract

from arXiv · show

This paper proposes a layer bit allocation method for Gemma-3-1B, formulating the problem as performance maximization (latency decrease) given a degradation budget constraint (allowable level of generation quality loss). This approach is different from time- and resource-consuming uniform layer quantization methods that are used in the literature (like GPTQ or AWQ) or allocation methods without proven performance-accelerating effect (like MixLLM or TorchAO). The layer sensitivity profile resulting from our prior work SA-PTQ is applied using the activation pass-through mode inside TensorRT-LLM. For each layer precision is determined individually in blocks, according to a grouping introduced in the prior step (5+5, 10+10, all26), differentiating the contribution of FFN, Attention, and lm_head to the overall speedup. The clock speed was measured for 13 W8A8 variants on an RTX 5090. We find that for FFN and lm_head the time cost of quantization/dequantization is compensated for by the use of integer arithmetic, while for short context lengths, the opposite holds true for Attention: an additional step of quantization slows execution down. We propose a manual implementation of SmoothQuant for TensorRT-LLM which was necessary due to export failures, unavailable for lm_head. The best solution found under joint consideration of all three criteria with minimal degradation was FFN 5+5 with lm_head, providing an 11.0% reduction in latency with negligible quality loss (98.90% Top-1 agreement, +0.85% perplexity degradation). With acceptable quality loss for FFN all26 + lm_head, a speedup up to 19.1% was found possible. We suggest further optimizations: fused attention kernels in INT8, KV-cache quantization, using FP8 instead of INT8 and partial Attention quantization analogous to FFN.

1 Introduction

The paper targets the unresolved computational trade-off in mixed-precision LLM quantization: layer sensitivities differ, but precision transitions add overhead. It proposes measuring and allocating layer bit-widths to maximize inference speed while controlling quality degradation.

  • Motivation: GPTQ and AWQ apply uniform bit-widths across layers, despite substantial differences in layer sensitivity to quantization.The prior sensitivity analysis reports FFN Q = 0.796 and tokenizer embedding Q = 0.819, with other layers between them.
  • Approach: The paper evaluates selective mixed-precision quantization using SA-PTQ sensitivity estimates and TensorRT-LLM activation pass-through mode.The approach measures per-layer computational effects rather than assuming mixed precision automatically accelerates inference.
  • Motivation: Mixed-precision inference incurs de-quantization and requantization overhead when activation tensors cross blocks with different precisions.This overhead can reduce or eliminate the computational benefit of lower-precision arithmetic.
  • Related work: MixLLM and TorchAO address mixed precision or quantization, but the cited literature does not establish which transformer layers benefit most from the approach or resolve transition overhead.MixLLM explicitly targeted memory consumption rather than inference speed, while uniform methods likewise do not solve mixed-precision overhead.
  • Approach: TensorRT-LLM provides an end-to-end activation-aware quantization path, but Q/DQ wrapper costs can exceed lower-precision arithmetic gains for some operations.The paper uses Gemma-3-1B experiments to measure this operation-dependent overhead.

2 Problem Statement

The study formulates Gemma-3-1B inference as maximizing tokens per second under an allowed quality-loss budget and profiles which subsystems offer useful acceleration. It prioritizes FFN and lm_head quantization while treating short-context Attention as a lower-priority target because its overhead and optimization complexity can limit gains.

  • Problem formulation: The optimization objective is to maximize tokens generated per second while keeping generation-quality degradation within a specified budget against FP16.
  • Profiling findings: 22.01 µs versus 13.03 µs versus 3.88 µs: FFN fused_fc, FFN proj, and Attention qkv have markedly different per-call costs.The profiling targets batch size 1 decode on an RTX 5090 using Nsight tools and TensorRT-LLM.
  • Profiling findings: 376.58 µs: lm_head dominates the profiled single invocation and exceeds the most expensive FFN sub-layer by almost an order of magnitude.Its high DRAM throughput makes lm_head a separate candidate for quantization alongside FFN.
  • Profiling findings: Attention is lower priority because improving its utilization requires more complex thread-count or memory-access restructuring than FFN or lm_head optimization.
  • Scope: The analysis is restricted to short contexts, while longer contexts with large accumulated KV-cache may substantially change Attention behavior.

3 FFN

This section applies SA-PTQ-guided blockwise W8A8 quantization to selected Gemma-3-1B FFN layers and evaluates speed, kernel behavior, and activation-quality trade-offs against FP16.

  • Method: SA-PTQ selects FFN layers using an SQNR-based sensitivity heatmap, where high SQNR indicates robustness and low SQNR indicates stronger quantization distortion.The reference output is yclean = x · Wfp16, while the quantized output is ydirty = x · Wint8.
  • Method: Symmetric per-channel quantization computes each output-channel scale from the maximum absolute weight value, using W8A8 with b = 8 and nearest-integer rounding.The same quantization procedure is applied to each selected linear FFN layer: gate_proj, up_proj, and down_proj.
  • Method: The 5+5 and 10+10 schemes quantize outer FFN layers, while unselected layers retain FP16 weights and activations.The 5+5 scheme covers layers 0–4 and 21–25; 10+10 covers layers 0–9 and 16–25.
  • Quality: 1.04× was the best speed-quality trade-off for 5+5, with activation SQNR above 9 dB and a maximum of 19 dB.The non-quantized middle layers partially compensate for errors from the first quantized layers; all26 instead shows a drop to 4–7 dB in layers 10–13.
  • Performance: 1.04–1.10× speedup was achieved by blockwise W8A8 FFN quantization at batch size 1 on the RTX 5090.TensorRT-LLM uses an INT8×INT8 tensor-core GEMM kernel for quantized layers, and increasing coverage increases its use on the RTX 5090.
  • Performance: Kernel choice is architecture-dependent: RTX 5090 reaches 1.10× with all26, whereas RTX 3070 favors 5+5 because broader coverage triggers the slower i8f32 kernel.On the RTX 5090, i8i8 use is not similarly limited by architecture.

4 Attention

This section tests W8A8 quantization of Attention projections alongside FFN quantization, finding no short-context batch-one speed gain while identifying profiling-based explanations and future optimization paths.

  • Method: Attention quantization targets the Q, K, V, and O linear projections within Gemma-3-1B decoder layers, in addition to the FFN quantization scope.The attn_ffn all26 configuration applies W8A8 to all 26 layers for both FFN and Attention projections.
  • Method: SmoothQuant shifts activation-range variation into per-channel-quantized weights using a per-channel scaling factor, making activations more uniform for per-tensor quantization.The calibration configuration uses α = 0.5, and the resulting scaling factor is stored as prequant_scaling_factor in the checkpoint patch.
  • Results: 1.069× was the overall speedup for attn_ffn all26, below the 1.090× achieved by FFN all26 alone at batch size 1 on the RTX 5090.The controlled comparison measured 407.28 tok/s for attn_ffn all26 versus 414.98 tok/s for FFN all26, a slowdown of −7.7 tok/s.
  • Future directions: At batch size > 1, experiments reached roughly 1350–1430 tok/s, corresponding to a ∼1.16× speedup, while FP8 FFN quantization reached 1.46× versus 1.10× for INT8.The paper also proposes KV-cache quantization and selective blockwise Attention quantization as future directions.
  • Mechanism: The masked_multihead_attention kernel remains FP16, so only the relatively small Q/K/V/O GEMMs are accelerated under Attention quantization.The GQA architecture has 1 KV head for 4 query heads, limiting the potential gain from quantizing these projections.

5 Attention + FFN

The section evaluates simultaneous Attention and FFN quantization, finding acceptable numerical quality but no realized speed benefit from the current TensorRT-LLM implementation. It attributes the gap to Q/DQ overhead, smaller attention projections, and unfused attention execution.

  • Results: Quantizing Attention alongside FFN produced acceptable quality but no real performance gain, with slight slowdowns versus FFN-only configurations.The paper separates numerical tolerance from hardware efficiency: attention precision remains admissible, but the current implementation does not exploit it efficiently.
  • Mechanism: Attention quantization wraps projections with explicit Q/DQ operations, breaking the continuous FP16 fused-attention path and adding overhead.The standard masked multi-head attention plugin is designed for a continuous FP16 path; INT8 insertion prevents that fused path.
  • Mechanism: Gemma-3-1B’s q_proj, k_proj, v_proj, and o_proj matrices are smaller than FFN matrices, so INT8 arithmetic gains are smaller while Q/DQ overhead remains comparable.Grouped-Query Attention reduces the absolute benefit available from quantizing the attention projections.
  • Further optimization: The current calibration uses per-channel SmoothQuant without specialized fused INT8-attention kernels, amounting to a naive Q/DQ integration.The paper identifies fused attention kernels, alternative calibration, and partial attention quantization as directions for recovering potential speed gains.

6 LM Head

The section treats lm_head separately because Gemma-3 ties it to the embedding matrix and standard TensorRT-LLM export fails for this component. A manual SmoothQuant implementation enables evaluation of its accuracy and speed.

  • Architecture: Gemma-3’s lm_head is tied to model.embed_tokens, so one shared weight tensor serves both input embedding lookup and output-logit computation.This shared architecture distinguishes lm_head from FFN and Attention weights, which are used once and are unrelated to other model components.
  • Implementation: Standard modelopt-to-TensorRT-LLM export reproducibly broke Gemma lm_head generation or silently omitted its quantization, with no official fix available during the experiment.The exporter could produce incoherent repetitive text or leave quant_algo null despite completed calibration.
  • Implementation: The authors therefore implemented SmoothQuant manually for lm_head, bypassing export_tensorrt_llm_checkpoint.They also resolved Gemma-specific checkpoint-export issues, including the required decoder_type value ’gemma3’.
  • Method: SmoothQuant redistributes activation outliers onto weights while preserving xW ⊤, enabling more uniform dynamic ranges before INT8 quantization.The chosen α = 0.5 splits the smoothing burden evenly between activations and weights.
  • Method: lm_head activations are quantized dynamically per token, whereas weights are quantized statically during checkpoint preparation.Per-token recomputation accommodates changing activation distributions during inference.
  • Results: 98.90% top-1 agreement was obtained on held-out next-token prediction, while isolated lm_head quantization achieved a 1.071× speedup over FP16.The speed measurement used the median of 15 runs with 50 generated tokens.

7 Combined Configurations

The combined-configuration experiments compare 13 subsystem combinations under a common wall-clock protocol. Results show that the strongest nominal speedups can coincide with severe quality degradation, while FFN 5+5 with lm_head offers the best combined outcome.

  • Experimental setup: 13 configurations were measured using the same wall-clock methodology on one RTX 5090, enabling direct comparison across results.Measurements used three warm-up runs, the median of 15 runs, 50 generated tokens, and a fixed prompt.
  • Speed–quality trade-off: 1.156–1.191× nominal speedups came from configurations including FFN 10+10, FFN all26, or full Attention quantization, but their text collapsed into repeated single tokens.The associated SQNR traces consistently fell below the 40 dB threshold.
  • Best combined configuration: The FFN 5+5 + lm_head configuration was the best result under combined speed and quality criteria.The supplied passage identifies this configuration as the best combined result, while its quantitative continuation is not included here.

8 Conclusion

The experiments support blockwise W8A8 quantization of Gemma-3-1B in TensorRT-LLM as a viable route to faster inference while preserving generation quality in selected configurations. The main remaining opportunity is more effective Attention engineering and further tooling support for lm_head quantization.

  • Conclusion: The experiments answer yes: blockwise W8A8 quantization of FFN, Attention, and lm_head can speed Gemma-3-1B inference while preserving generation quality.The conclusion qualifies this answer with an important nuance that motivates future directions.
  • FFN and lm_head: FFN quantization provides predictable speed benefits, with 1.031× for 5+5 and 1.075× for 10+10, while quality drops more noticeably at broader coverage.The passage reports only slight quality decline for 5+5 and noticeable decline for 10+10 and all26.
  • Attention: Adding INT8 Attention quantization reduces throughput relative to FFN-only quantization at every tested FFN coverage level.Reported comparisons are 1.023× versus 1.031×, 1.054× versus 1.075×, and 1.069× versus 1.090×.
  • Future directions: Nsight Systems analysis attributes the Attention trade-off to FP16 execution of masked_multihead_attention and overhead that is not offset by quantizing q/k/v/o projections.The passage identifies fused attention kernels as the key engineering direction.
  • Overall outcome: The study reports 11.0% lower inference time with minimal quality loss for FFN 5+5 + lm_head and up to 19.1% lower time with noticeable quality loss for FFN all26 + lm_head.The conclusion frames these as the practical benefits and quality trade-off of blockwise W8A8 quantization.
Loading 2608.28003v1…