PaperlayerView exact revision on arXiv

Source-linked AI summary

EFQ-Softmax: Exp-Free Quantization for Softmax

Haohui Han, Yuming Wan, Hongni Wang, Pengcheng Xie, Xiaodong Yan, Runqi You, Wencong Zhang

arXiv:2609.09721v1cs.LG

TL;DR

Low-bit attention still uses a high-precision exp-then-quantize probability path, creating a mismatch with its low-bit matrix consumer. EFQ-Softmax directly generates block-scaled E2M1 probabilities from shifted scores and uses them in both online-attention updates. It improves benchmark means, reduces A5 vector-stage latency by 40.33%, and preserves reported end-to-end quality.

  • Problem

    Low-bit attention commonly produces shifted-score exponentials and temporary high-precision probabilities before quantizing them for low-bit PV multiplication, mismatching the probability producer and consumer.

  • Method

    EFQ-Softmax directly maps shifted attention scores to block-scaled E2M1 probability operands using exponent-only local scaling, residual normalization, and affine code generation.

  • Results

    EFQ-Softmax raises Qwen3-8B's seven-task mean from 0.6749 to 0.6773, Qwen3-VL's nine-task mean from 0.7826 to 0.8000, and reduces A5 vector-stage latency by 40.33% on average.

  • Takeaways & Limitations

    Direct low-bit probability generation can replace post-softmax probability quantization while preserving the numerical behavior and end-to-end quality required by online attention.

Abstract

from arXiv · show

Low-bit attention accelerates Transformer inference by moving the $QK^\top$ and $PV$ matrix multiplications to FP8 or FP4 matrix engines. However, the softmax path often evaluates shifted-score exponentials in higher precision, forms a temporary probability block, and quantizes it before low-bit $PV$ multiplication. This exp-then-quantize path creates a mismatch between a high-precision probability producer and a low-bit matrix consumer. We propose EFQ-Softmax (Exp-Free Quantization for Softmax), a low-bit probability-generation method that directly maps shifted attention scores to block-scaled E2M1 operands. For each microscaling block, EFQ-Softmax selects an exponent-only scale from the local maximum, maps the shifted scores to a normalized residual domain, and generates nonnegative E2M1 probability codes using a single affine rule. The resulting operand is used consistently in both the $\widetilde{P}V$ numerator update and the $\widetilde{P}\mathbf{1}$ denominator update. The FlashAttention-style row-maximum update, historical rescaling, high-precision accumulation, and final normalization remain unchanged. We evaluate end-to-end quality on Qwen3-8B, Qwen3-VL-8B-Instruct, and WAN2.2-TI2V-5B, and separately measure kernel-level performance on the A5 vector unit. EFQ-Softmax improves the Qwen3-8B seven-task mean from 0.6749 with MXFP4 to 0.6773 and the Qwen3-VL nine-task mean from 0.7826 to 0.8000. On WAN2.2, it maintains temporal consistency and visual quality comparable to the FP16 and MXFP4 baselines under VBench. On the A5 vector unit, EFQ-Softmax reduces the vector-stage latency of the fused probability-generation kernel by 40.33% on average across sequence lengths from 16K to 128K. These results show that direct low-bit probability generation can replace the conventional exp-then-quantize path while preserving end-to-end model quality.

I. INTRODUCTION

Attention becomes increasingly costly at long sequence lengths because score and probability processing remain expensive even when tiled low-bit matrix multiplication accelerates QK^T and PV. EFQ-Softmax directly generates low-bit probability operands from shifted scores, preserving online-attention state updates while improving quality and probability-generation latency.

  • I. INTRODUCTION: FlashAttention reduces memory traffic through tiled online recurrence, but probability generation becomes a larger execution-time share as low-precision matrix multiplications accelerate.The conventional path produces probabilities in FP16 or FP32 before quantizing them for low-bit PV multiplication.
  • I. INTRODUCTION: The method removes dense current-block exponentials, temporary high-precision probability tiles, and separate post-exponential quantization from the conventional exp-then-quantize path.The row-level exponential for historical rescaling remains part of online attention.
  • I. INTRODUCTION: EFQ-Softmax directly converts shifted attention scores into block-scaled E2M1 probability operands through exponent-only scaling, residual normalization, and affine code generation.The generated codes and block-level E8M0 scale form the MXFP4 probability operand.
  • I. INTRODUCTION: The same generated probability operand feeds both numerator PV and denominator P1 updates while accumulation and final normalization remain high precision.Row-maximum updates, historical rescaling, and online recurrence state are retained.
  • I. INTRODUCTION: 40.33% lower A5 vector-stage latency is achieved across 16K–128K sequence lengths while end-to-end quality is preserved across language, vision-language, and video workloads.Qwen3-8B rises from 0.6749 to 0.6773, Qwen3-VL rises from 0.7826 to 0.8000, and WAN2.2 remains comparable under VBench.

B. Block-Scaled E2M1 Probability Representation

EFQ-Softmax represents each probability block with an exponent-only scale and nonnegative E2M1 codes. Local residual normalization separates block dynamic-range selection from elementwise code generation.

  • E2M1 representation: The block-scaled representation writes each probability as an E2M1 code multiplied by its shared scale.The scale determines exponent range, while the code supplies the local mantissa-like value.
  • Block scale: Exponent-only scaling selects a shared power-of-two scale from each microscaling block’s local maximum.The scale is reused by all block elements, avoiding general floating-point scale generation.
  • Residual normalization: Residual log-domain normalization converts each shifted score into a locally bounded input for code generation.Using the normalized residual avoids mixing block-level range selection with elementwise code selection.
  • Microscaling blocks: Microscaling block size trades approximation error against scale metadata and scale-generation work.Smaller blocks provide more local scaling, while larger blocks reduce overhead.

C. EFQ-Softmax Probability Code Generation

EFQ-Softmax generates the final low-bit probability operand directly from shifted scores using affine residual thresholds. This removes the intermediate high-precision probability tile and explicit current-block exponential.

  • Procedure overview: Figure 3 summarizes five stages that transform shifted attention scores into packed block-scaled E2M1 operands.The figure presents the complete probability-generation procedure.
  • Direct code generation: Direct affine code generation produces the final packed E2M1 probability operand without materializing an intermediate FP16/FP32 probability block.The output of the affine rule is already the 4-bit operand consumed by the attention data path.
  • Code mapping: The affine rule clips residuals into the nonnegative E2M1 range, with saturation yielding codes 0 through 7.Code 7 represents the largest nonnegative E2M1 value, 6.
  • Threshold design: EFQ-Softmax replaces scalar nearest-rounding after explicit exponentiation with uniformly spaced residual thresholds.This trades scalar nearest-rounding optimality for a simpler probability-code generation path.
  • Calibration and target: The method targets output-level attention accuracy rather than exact scalar exponential reconstruction.The affine parameters τ and h are selected offline through calibration.

D. EFQ-Softmax-LUT Variant

EFQ-Softmax-LUT retains the same block scale and E2M1 output set but inserts a 16-level intermediate mapping. It serves as an accuracy-oriented ablation of mapping flexibility.

  • Two-stage mapping: EFQ-Softmax-LUT adds a 16-to-8 remapping stage between residual indexing and the final E2M1 code.Its folding table maps 16 intermediate indices to the eight-code output set.
  • Ablation role: The variant uses finer intermediate boundary placement than the eight-code EFQ-Softmax rule.It is intended to measure accuracy gains from increased mapping flexibility at probability-code generation.
  • Mapping comparison: Figure 4 compares EFQ-Softmax’s affine mapping with EFQ-Softmax-LUT’s two-stage mapping.The comparison isolates the difference in code-mapping structure.

A. Integration with Online Attention

EFQ-Softmax integrates into FlashAttention-style online attention by replacing only current-block probability generation. The same generated operand feeds numerator and denominator updates while row statistics and accumulation remain unchanged.

  • Numerator update: The generated low-bit probability operand is consumed by the numerator branch for the subsequent probability-value multiplication.The value-side representation is an existing low-bit path; EFQ-Softmax contributes the probability operand.
  • Denominator update: The same probability operand feeds the denominator row-sum branch, whose output remains accumulated in FP16 or FP32.This keeps both online branches based on one probability approximation.
  • Preserved recurrence: The FlashAttention online structure preserves row-maximum updates, historical rescaling, high-precision accumulation, and final division.EFQ-Softmax modifies only current-block generation rather than the surrounding recurrence.
  • Current-block path: EFQ-Softmax removes dense current-block exponentials and replaces them with affine index generation.The row-level historical rescaling exponential remains part of the online recurrence.

B. Algorithm Summary

EFQ-Softmax uses offline calibration to select global threshold parameters, then generates each microscaling block’s E2M1 probability codes through a calibrated affine rule within online attention.

  • Parameter calibration: Offline calibration selects (τ, h) by minimizing attention-output error over representative score blocks, with the chosen operating point fixed during inference.The objective captures coupled effects on both numerator and denominator updates rather than scalar exponential error alone.
  • Block processing: For each microscaling block, the algorithm finds the local maximum, derives an exponent-only power-of-two scale, and normalizes scores into a residual domain.The residual-domain input remains bounded after block-scale normalization.
  • Code generation: A single affine threshold rule maps normalized residuals to nonnegative E2M1 probability codes, which are packed with the block scale into the final operand.The parameters τ and h shift and space the affine thresholds, respectively.
  • Online update: The generated probability operand is consumed by the online update while preserving the existing state-update structure and high-precision accumulation.Calibration is performed offline rather than adapting parameters per input.

D. Algorithmic Cost and Method Boundary

EFQ-Softmax removes dense current-block exponentials and intermediate high-precision probability materialization, placing direct score-to-code generation immediately after score shifting and before both attention branches.

  • Removed operations: The method replaces dense exponentials, temporary high-precision probability formation, and subsequent quantization with residual-domain score-to-code generation.Per-element work uses subtraction, multiplication, flooring, clipping, and bit-pattern generation.
  • Kernel boundary: EFQ-Softmax receives shifted score blocks, microscaling partitions, and fixed parameters, then returns the low-bit probability operand.Masks are applied before score shifting and probability generation.
  • Kernel boundary: The generated code and scale are consumed by both numerator and denominator updates immediately after the updated row-maximum shift.The value-side representation and online accumulator precision remain unchanged.

A. Experimental Setup

Experiments replace self-attention and cross-attention with quantized attention across language, vision-language, and video models, comparing FP16, MXFP4, and calibrated EFQ-Softmax variants.

  • Experimental setup: The evaluation covers seven zero-shot language tasks, nine vision-language tasks, and fixed-setting text-to-video workloads using calibrated EFQ-Softmax operating points.The Qwen3-8B and Qwen3-VL comparisons use task-averaged metrics alongside per-task results.
  • Qwen3-8B: Batch-size replication keeps each Qwen3-8B method within 0.006 of its batch-8 seven-task mean, preserving the comparison’s conclusion across batch sizes 8, 16, and 24.FP16 remains highest and MXFP4 lowest, while EFQ-Softmax variants occupy the intermediate range.
  • Qwen3-VL-8B-Instruct: Qwen3-VL EFQ-Softmax-Balance raises the 9-task mean from 0.7826 with MXFP4 to 0.8000, outperforming MXFP4 on all six grounding splits.VQA performance remains nearly indistinguishable, with a maximum drop of 0.0032 and a TextVQA gain of 0.0007.

D. Text-to-Video Generation on WAN2.2-TI2V-5B

On WAN2.2, EFQ-Softmax preserves temporal and visual quality relative to FP16 and MXFP4 under matched generation settings, while its fused probability kernel substantially reduces vector-stage latency.

  • VBench quantitative metrics: EFQ-Softmax achieves temporal-consistency scores comparable to FP16 and MXFP4, including subject consistency 0.9493 and motion smoothness 0.9891.It also records the highest imaging quality at 0.7178 versus 0.7098 for FP16 and 0.7004 for MXFP4 SCALE4.
  • Visual quality: Matched 14-prompt video comparisons show EFQ-Softmax retaining FP16-like motion coherence, lighting, and fine texture details while remaining comparable to MXFP4 SCALE4.All variants use identical prompts, generation parameters, and random seeds.
  • Kernel performance: The fused EFQ-Softmax LUT reduces A5 vector-stage latency by an average of 40.33% across sequence lengths from 16K to 128K.The reduction remains nearly invariant, ranging from 40.10% at 16K to 40.48% at 128K.
  • Method boundary: EFQ-Softmax directly generates the low-bit probability operand and feeds it to both numerator and denominator branches while retaining normalized online attention accumulation.The method changes current-weight generation without changing final normalization.
Loading 2609.09721v1…