VC-Attention: Value Smoothing and Softmax Casting for Low-bit Attention

AI summary

Overview

Research area: Efficient inference for video diffusion transformers — specifically low-bit quantization of the attention mechanism.

Technical level: Advanced. The paper assumes familiarity with FlashAttention-style online softmax, floating-point formats (FP8/E4M3, NVFP4), Tensor Core pipelines, and diffusion transformer inference.

Scope: A training-free low-bit attention kernel that removes two deployment bottlenecks in video DiTs — value outliers that dominate output error and the FP32 exponential/cast that dominates softmax latency — implemented and benchmarked across four video models and five GPUs.

What This Paper Is About

Video diffusion transformers (DiTs) spend most of their inference time in attention, and low-bit Tensor Cores promise large speedups — but only if the quantized attention is both accurate and actually faster. Accuracy collapses because a few outlier values set the quantization scale for entire blocks of tokens; speed collapses because the FP32 exponential and format-cast inside softmax become the longest pipeline stage once the matrix multiplications run in low precision. VC-Attention attacks both problems jointly with a value-side smoothing technique and a direct probability encoding, without any retraining or per-model calibration.

Key Contributions

1. V-Smooth: An online, value-guided token reordering. Tokens are grouped by a lightweight online k-means on the value tensor, keys and values are permuted together by the resulting label ordering, and each hardware block then subtracts its mean and quantizes only the residual. The mean is restored inside the existing online-softmax recurrence using the row sum the kernel already accumulates, so no extra pass or buffer is required.

2. ExpCast-FP8: A fused probability cast that writes the E4M3 byte directly from the log-domain score with a single fused multiply-add, replacing both the FP32 exponential and the FP32-to-FP8 conversion. The construction exploits the fact that an E4M3 byte is an affine function of the log of the value it encodes.

3. A fused CuTe/CUDA implementation that absorbs the entire preprocessing chain (rotary embedding, means, gather by the permutation, QK Hadamard, quantizers) so permuted and rotated high-precision tensors never touch HBM, plus an amortized grouping schedule that runs grouping only on the first quarter of denoising steps.

4. A row-level error guarantee (Proposition 3.1) bounding the total-variation distance between exact and ExpCast-FP8 probability vectors at 3.64% plus an underflow tail, with a matching bound on the output vector in terms of the diameter of the value set.

Main Findings

• Value error, not score error, bounds fidelity on video DiTs. After smoothing queries and keys or applying a Hadamard rotation, the value term of the output error accounts for 82% of it on Wan2.2. A Hadamard rotation of QK shifts PSNR by at most 0.1–0.3 dB across all models, confirming that the removable score-product error is already gone.

• Sorting is what makes block demeaning useful. Over 100 Wan2.2 heads, the block mean removes 8% of block energy in sequence order, 12% under a static spatiotemporal cube, and 36% after k-means grouping. A Hadamard rotation of V changes the value error by only 0.2%; the fixed cube recovers 3.7%.

• V-Smooth is the most faithful training-free low-bit attention at both precisions. At 8 bits it improves PSNR over SageAttention2 by 2.3 dB on Wan2.2 and 2.8 dB on HunyuanVideo-1.5, reducing LPIPS by 13–29%. At 4 bits it beats SageAttention3 by 2.9 dB on Wan2.2 and 3.6 dB on LongCat-Video, with LPIPS reductions up to 41%.

• ExpCast-FP8 buys speed at a modest fidelity cost. Fusing it trades 0.7–2.1 dB of PSNR for the reported latency gains, and the combined kernel still outperforms SageAttention2 on all four models.

• Speedups convert to real wall-clock gains. Attention runs 1.59× faster than BF16 FlashAttention-4 on B200 (6.02× over SageAttention2) and 1.46× on H200. At 4 bits on workstation Blackwell, V-Smooth reaches 2.27× on the RTX PRO 6000 and 3.58× on the RTX 5090. End-to-end clip generation is 1.13–1.19× faster on datacenter GPUs and 1.36–1.70× faster on workstation cards.

• Grouping overhead is amortized to 3–4% of attention time. Grouping on the first quarter of denoising steps reaches 26.4 dB PSNR versus 23.8 dB when the same number of grouping steps is spread uniformly, and 26.9 dB when grouping on every step — so where the grouping steps sit matters more than how many there are.

• Training-free Attn-QAT is not competitive. Run without its model-specific retraining, Attn-QAT falls 3.4–6.7 dB of PSNR below SageAttention2 on every model, and is the only method that perturbs the VBench subject-consistency and imaging-quality scores.

Methodology in Plain English

The researchers start from the observation that low-bit attention has two separable failure modes and treat them independently.

For accuracy, they note that a quantizer assigns one shared scale to a block of values, so a single large entry in that block shrinks the effective resolution available to every other entry. Value outliers in video DiTs are spread unpredictably across tokens, heads, layers, and denoising steps — they follow no fixed channel or spatial pattern, which rules out rotations and static layouts. The fix is to change which tokens share a block. A cheap online k-means over the value tokens assigns each token a cluster label; sorting those labels produces a permutation that is applied to both keys and values. Because permuting keys permutes the columns of the probability matrix exactly as it permutes the rows of the value matrix, the attention output is mathematically unchanged. Tokens in the same block now share a common component, so subtracting the block mean removes a much larger share of the energy, and only the residual is quantized. The mean is added back during the online softmax recurrence as a rank-one outer product using the row sum the kernel already tracks.

For speed, they observe that an E4M3 floating-point byte is by construction a logarithmic encoding: its exponent field is the integer part of log₂ of the value plus a bias, and its mantissa field approximates the fractional part. That makes the byte an affine function of log₂ of the value. Since online softmax already keeps every score in the log domain after subtracting the running row maximum, the byte can be produced with one multiply-add and one integer conversion instead of an FP32 exponential followed by a cast. The constant term is the minimax centering of the one residual error term, not a fitted parameter; the clip range keeps codes in the normal E4M3 range and maps underflow to zero. The result is reinterpreted as an E4M3 value rather than converted, which costs no instruction.

The kernel is hand-written in CuTe/CUDA by modifying FlashAttention-4/SageAttention in place, fusing the preprocessing chain backwards from the quantizer so intermediate tensors stay on chip, and restricting grouping to the first quarter of denoising steps, where reusing the permutation across adjacent steps is safe because attention layouts change little.

Why This Matters

Impact on research. The paper reframes low-bit attention for video generation: the field has concentrated on the score product (SageAttention, FlashAttention-3's Hadamard path), but this work shows the value operand is the binding constraint on fidelity for video DiTs, and that softmax's scalar path is the binding constraint on latency for datacenter GPUs. It also provides a rare example of an exact, constant-free transformation from a log-domain score to an FP8 code, which may transfer to other softmax-based kernels.

Real-world applications:

• Video generation serving. Production text-to-video services can cut per-clip GPU cost by 13–19% on datacenter hardware without retraining or changing the sampling schedule.
• Workstation and prosumer deployment. The 4-bit path yields 1.36–1.70× end-to-end speedups on RTX PRO 6000 and RTX 5090, which matters for creative studios running models locally.
• Long-context and high-resolution video. The speedups scale with sequence length, so the technique is most valuable exactly where attention dominates — 70K-token 720p clips and beyond.
• Composable acceleration. Because VC-Attention leaves the attention pattern and denoising schedule untouched, it stacks with sparse attention, distillation, feature caching, and sequence parallelism.

Industry relevance. The results are measured on shipping datacenter (B200, B300, H200) and workstation (RTX PRO 6000, RTX 5090) hardware, against released baselines including SageAttention2/3 and FlashAttention-4. The training-free property is the key commercial point: no model-specific data collection, no fine-tuning runs, no per-model calibration, and the kernel drops into an existing FlashAttention-derived pipeline.

Future Directions

• Extending ExpCast to 4-bit datacenter configurations. The paper explicitly does not evaluate a 4-bit datacenter setup because an on-the-fly NVFP4 probability places its per-16-element scale computation on the softmax critical path, which ExpCast-FP8 does not remove. Eliminating that scale computation is the natural next step.

• Replacing or accelerating the grouping itself. Plain k-means was chosen over balanced k-means because balancing costs 2.8× more for an 8.5% error improvement. A cheaper clustering that stays near the balanced operating point would widen the margin over baselines.

• Do the mean-restoration trick and permutation generalize beyond video DiTs? The rank-one mean correction inside the online recurrence is a general identity. Whether it helps long-context language models, where value outlier structure differs, is untested.

• Interaction with sparse and linear attention. The paper argues the method composes with sparsity and linear attention, but provides no joint measurements. Combining value-side quantization with token-skipping raises questions about whether skipped tokens should participate in the clustering.

Target Audience

This paper is aimed at systems and inference-optimization researchers working on diffusion transformers and video generation, GPU kernel engineers interested in FP8/NVFP4 numerical formats and FlashAttention-derived pipelines, and practitioners deploying large video models who need lower serving cost without retraining. Readers will benefit most from prior familiarity with online softmax, block quantization, and Tensor Core precision formats; readers without that background can still follow the problem framing and the headline fidelity and speed results.

Authors’ abstract

Diffusion Transformers deliver state-of-the-art video generation, but their long spatiotemporal sequences make attention the dominant deployment cost, and a deployable low-bit kernel must be accurate and fast. Accuracy is limited by outliers: a block's quantization scale is set by its largest entries, leaving typical entries confined to a narrow range of representable values. Prior work smooths queries and keys, but value outliers follow no fixed channel or spatiotemporal structure and remain the dominant source of output error. Speed is limited by softmax: low-bit Tensor Cores accelerate only the two matrix multiplications, so the high-precision exponential between them becomes the longest pipeline stage on datacenter GPUs. We propose VC-Attention, a training-free low-bit attention framework that addresses both by pairing Value smoothing with a fused probability Cast. V-Smooth reorders value tokens by lightweight online clustering, so the tokens in a hardware block quantize well together. It quantizes only the residual after subtracting the block mean, and restores that mean from the row sum the online softmax already maintains. ExpCast-FP8 maps log-domain scores directly to E4M3 probability codes with one fused multiply-add, eliminating the FP32 exponential and the format conversion. We implement VC-Attention for B200, B300, H200, RTX PRO 6000, and RTX 5090. Across Wan2.2, LongCat-Video, HunyuanVideo-1.5, and MiniMax-H3, VC-Attention improves fidelity over low-bit baselines, speeds up the attention kernel over BF16 FlashAttention-4 by 1.46-1.59x on datacenter Blackwell and Hopper and by 2.3-3.6x on workstation cards, and generates a clip 1.13-1.19x and 1.36-1.70x faster end to end.