CDFQuant

CDFQuant is a simplex-preserving post-softmax quantization method for Transformer attention. Instead of quantizing attention probabilities directly, it quantizes cumulative distribution values and reconstructs probabilities by differencing, which preserves non-negativity and the unit-sum constraint without renormalization.

This repository contains the reference implementation used to evaluate CDFQuant on ImageNet-style ViT models and AST on Speech Commands V2, together with small benchmarking utilities and optional baseline adapters.

Repository Scope

This public repository intentionally excludes:

• datasets
• pretrained checkpoints
• experiment outputs and logs
• vendored third-party baseline repositories

Those assets stay local and are ignored by Git so the public repo remains lightweight and reviewable.

Contents

• cdfquant.py: core quantizers, E_max computation, and timm attention patching
• eval.py: ImageNet-style evaluation for timm ViT backbones
• eval_ast.py: AST evaluation on Speech Commands V2
• run_ast_baseline.py: batch runner for AST baselines
• bench_row_ops.py: micro-benchmark for row-wise quantization kernels
• run_power_eval.py: wrapper that samples nvidia-smi power/utilization during evaluation
• eval_repq_emax.py: optional RepQ-ViT evaluation helper
• eval_fqvit.py: optional FQ-ViT evaluation helper
• eval_fqvit_lis.py: optional LIS-only baseline helper
• eval_aphq_insert.py: optional APHQ-ViT insertion study helper
• tests/: lightweight unit tests for the core quantizer behavior

Environment

The scripts were written for Python 3.10+ with PyTorch-based tooling.

Install runtime dependencies:

pip install -r requirements.txt

Install test dependencies:

pip install -r requirements-dev.txt

Quick Start

Run unit tests:

pytest

Evaluate CDFQuant on a timm ViT model:

python eval.py \
  --model deit_small_patch16_224 \
  --data /path/to/imagenet/val \
  --checkpoint /path/to/model.safetensors \
  --mode cdfquant \
  --bits 4

Evaluate AST on Speech Commands V2:

python eval_ast.py \
  --mode cdfquant \
  --bits 4 \
  --data-dir /path/to/speech_commands_v2 \
  --checkpoint /path/to/ast_speechcommands.pth \
  --ast-model-file /path/to/ast_repo/src/models/ast_models.py

Run the row-level micro-benchmark:

python bench_row_ops.py \
  --bits 4 \
  --n 197 \
  --batch 2048 \
  --heads 12 \
  --output-prefix artifacts/bench_row_ops

Optional External Baselines

The baseline helper scripts do not vendor third-party code. By default they look for sibling directories such as FQ-ViT, RepQ-ViT, and APHQ-ViT. You can also point them to local clones with environment variables:

• FQVIT_ROOT
• REPQVIT_ROOT
• APHQ_VIT_ROOT

Examples:

export FQVIT_ROOT=/path/to/FQ-ViT
export REPQVIT_ROOT=/path/to/RepQ-ViT
export APHQ_VIT_ROOT=/path/to/APHQ-ViT

Notes

• eval.py and eval_ast.py are the main entry points for the CDFQuant method itself.
• The repository ignores generated artifacts by default; create local output directories as needed.
• run_power_eval.py requires nvidia-smi to be available in PATH.

Citation

If you use this repository, please cite the accompanying CDFQuant paper and this code release. A machine-readable citation file is provided in CITATION.cff.

License

This repository is released under the MIT License. Third-party baselines used for comparison remain under their own original licenses and are not redistributed here.