⚡ VAR-Q: KV-Cache Quantization for Visual Autoregressive Generation

VAR-Q is a lightweight KV-cache quantization method for efficient visual autoregressive generation. It reduces inference-time KV-cache memory while preserving generation quality, and is designed to attach to existing visual autoregressive model implementations without modifying their source code.

⚡ Quick Start in 60 Seconds

This smoke path does not require third-party backends, checkpoints, or datasets.

git clone https://github.com/lujiaji/VAR-Q.git
cd VAR-Q
git checkout varq_official
# Install PyTorch first from https://pytorch.org/get-started/locally/
pip install -e .
python examples/smoke_quant_roundtrip.py
python examples/smoke_pack_unpack.py
python examples/smoke_hook_mock_attention.py

For development checks:

pip install -e ".[dev]"
pytest tests/

🎯 Scope

VAR-Q is a runtime KV-cache quantization method for visual autoregressive generation. It quantizes K/V cache tensors created during inference; it is not a weight-only quantization method.

Weight quantization methods such as GPTQ, AWQ, and related approaches are orthogonal to VAR-Q. In principle, a backend can combine weight-only quantization for model parameters with VAR-Q for runtime KV-cache memory.

✨ Highlights

• Runtime hook integration: VAR, Infinity, InfinityStar, LiveTalk, Self-Forcing, and LongLive are instrumented in memory after model construction; no forked backend source is distributed.
• Main VAR-Q method: supports VARQ, all G_* grouping variants, ratio-controlled grouping, pre-RoPE control, and low-bit KV cache packing/unpacking.
• Memory-conscious runtime: stores compact scale metadata, avoids token-expanded scale caches, supports exact cache preallocation, and uses compiled CUDA quant/pack/unpack/dequant operators on CUDA.
• Clean ablation boundary: KIVI, FLexGen, sparse-attention, and GPTQ W4/A8 reference implementations live under ablation/ and are installed as an optional comparison package alongside VAR_Q.
• Backend-friendly release: third-party model repositories, checkpoints, generated media, and experiment scratch files are ignored by default.
• Minimal core dependency: Python import and CPU reference execution only depend on PyTorch; CUDA inference uses the bundled CUDA extension. Backend-specific environments should follow the upstream model repositories.

📊 Deployment Results

Reported A100-80GB evaluations show approximately 75% KV-cache reduction at INT4 and up to 87% at INT2. The reduced cache enables substantially larger full-generation batches across image and video backends.

For Infinity-8B, the effective compression ratios below include scale metadata:

KV-cache reduction refers to persistent K/V storage rather than whole-process peak GPU allocation. Maximum batch is the largest completed full-generation run under each reported precision.

Fused-kernel overhead reduction

On Infinity-8B at batch size 1, fusing packed-KV loading, unpacking, and dequantization into attention reduces both attention-path and end-to-end latency relative to the same quantized runtime without fusion:

The speedups above compare fused and unfused quantized execution. They do not claim that quantized execution is faster than the BF16 reference (3164 ms E2E).

🧩 Supported Backends

All listed entrypoints install hooks at runtime and do not require a committed patch to their upstream repositories. Standard Wan2.1/Wan2.2 diffusion is not listed because it does not maintain a persistent autoregressive KV cache.

For backends using the general hook router, VAR-Q and ablation methods are selected through the same attachment point while their quantization implementations remain isolated: VAR-Q is implemented in VAR_Q/, and comparisons remain in ablation/.

📦 Repository Layout

VAR-Q/
├── VAR_Q/                  # VAR-Q core and runtime hooks
│   ├── quant.py            # VARQ, G_* grouping, ratio logic, pre-RoPE control
│   ├── pack_unpack.py      # Low-bit pack/unpack utilities
│   ├── config_loader.py    # Public JSON normalization
│   ├── paths.py            # Relative third-party path helpers
│   ├── fused/              # Packed-KV attention Python API
│   ├── csrc/               # CUDA quant/pack/unpack/attention extension
│   └── hooks/              # Model hook installers
├── ablation/               # KIVI / FLexGen / KVQuant implementations
├── Benchmark/              # Public evaluation entrypoints
├── configs/                # Curated JSON configs
├── docs/                   # Public documentation
├── examples/               # No-checkpoint smoke demos
├── scripts/                # Public inference and evaluation launchers
├── tests/                  # Lightweight public smoke tests
├── third_party/README.md   # Upstream checkout instructions
├── requirements-varq.txt   # Lightweight VAR-Q convenience environment
├── pyproject.toml          # Editable install metadata
└── LICENSE

Ignored local-only paths include third_party/, temp/, benchmark outputs, model weights, generated media, and test output/tmp directories. The lightweight tests/ suite is part of the public repository.

🚀 Installation

Clone VAR-Q and check out the public branch:

git clone https://github.com/lujiaji/VAR-Q.git
cd VAR-Q
git checkout varq_official

Install PyTorch first, matching your CUDA driver and runtime. Follow the official PyTorch selector rather than relying on VAR-Q to pin a wheel:

# Example only. Choose the command for your CUDA version from:
# https://pytorch.org/get-started/locally/
pip install torch torchvision --index-url https://download.pytorch.org/whl/cu121

Install VAR-Q in editable mode:

conda activate <your-backend-env>
pip install -e .

Build the bundled CUDA extension before CUDA inference:

bash scripts/bench/build_fused_flash.sh

import VAR_Q and the CPU reference path remain available on CPU-only hosts; CUDA tensors never fall back to Triton or a Python pack/unpack implementation.

For a non-editable minimal convenience install after PyTorch is already installed, requirements-varq.txt is retained:

pip install -r requirements-varq.txt

VAR-Q itself only needs PyTorch for CPU import/reference use and a compatible CUDA toolchain for CUDA inference. VAR, Infinity, InfinityStar, Self-Forcing, and LongLive may pin different CUDA, PyTorch, flash-attn, xformers, tokenizer, or evaluation package versions; install those dependencies inside each official upstream repository instead of putting them into the VAR-Q root environment metadata.

🧪 Tested Environments

🧱 Third-Party Models

Clone upstream repositories into third_party/:

git clone https://github.com/FoundationVision/VAR third_party/VAR
git clone https://github.com/FoundationVision/Infinity third_party/Infinity
git clone https://github.com/FoundationVision/InfinityStar third_party/InfinityStar
git clone https://github.com/ChenhongyiYang/LiveTalk third_party/LiveTalk
git clone https://github.com/guandeh17/Self-Forcing third_party/Self-Forcing
git clone https://github.com/NVlabs/LongLive third_party/LongLive

VAR-Q does not vendor third-party source or distribute patch files. Public launchers load the backend model and install hooks in memory.

Checkpoints are intentionally not stored in JSON configs. Provide them through command-line arguments, environment variables, or the upstream backend's native loader.

Third-party backends, checkpoints, datasets, and generated assets are not distributed with this repository. They are governed by their own licenses and usage terms.

🪝 Runtime Hook API

VAR-Q follows a SmoothQuant-style runtime replacement design. The installer scans supported attention modules, stores the original methods in a handle, and replaces the instance-level KV-cache path in memory.

from VAR_Q.hooks import install_varq_hooks, remove_varq_hooks

handle = install_varq_hooks(
    model,
    model_type="var",  # "var", "infinity", or "infinitystar"
    quant_config={
        "enable": True,
        "q_bits": 4,
        "quant_method": "VARQ",
        "qkv_format": "BLHc",
        "pack_to_int32": True,
    },
)

# Run backend inference.

remove_varq_hooks(handle)

Use quant_method="VARQ" or any G_* method for the main method. Use KIVI, FLexGen, or KVQuant in public configs for ablations; the loader normalizes them to isolated ablation implementations.

LiveTalk uses a dedicated hook installer for its chunk-overwrite cache semantics:

from VAR_Q.hooks import install_livetalk_hooks, remove_livetalk_hooks

handle = install_livetalk_hooks(pipeline, quant_config)
# Run LiveTalk generation.
remove_livetalk_hooks(handle)

Self-Forcing and LongLive expose direct installers. Their launchers call these automatically after the upstream pipeline has been constructed:

from VAR_Q.hooks import install_longlive_hooks, install_self_forcing_hooks

sf_handle = install_self_forcing_hooks(sf_pipeline, "configs/self_forcing/varq/base/SF-VARQ-4.json")
ll_handle = install_longlive_hooks(ll_pipeline, "configs/longlive/varq/base/LL-VARQ-4.json")

The underlying VideoKVCacheAdapter remains available as an integration API for additional causal video models.

🧠 Memory-Efficient Runtime

VAR-Q reduces active cache allocations in addition to reporting packed byte counts:

• Scale metadata is stored compactly per group rather than expanded over every token.
• Packed K/V buffers can be preallocated with expected_total_seq_len and preallocate_kv_cache when the generation length is known.
• quant_compute_dtype="native" and dequant_dtype="native" avoid unnecessary full-size FP32 or cast temporaries.
• CUDA uses compiled quantize+pack, pack, unpack, unpack+dequant, and packed-KV attention operators; CPU execution retains a small PyTorch reference path.
• dequant_workspace_policy="release" avoids retaining dense dequant workspaces between attention calls.

For memory measurements, enable the entrypoint's profiling option where available. Reports include packed K/V bytes, scale bytes, dequant workspace bytes, and PyTorch allocated/reserved peaks. Compare backends using identical prompts, generation schedules, batch size, dtype, and a clean CUDA device.

🎯 Inference Scripts

All public scripts resolve paths relative to the VAR-Q repository root.

VAR:

export VARQ_VAE_CKPT=/path/to/vae_ch160v4096z32.pth
export VARQ_VAR_CKPT_TEMPLATE='/path/to/var_d{}.pth'
bash scripts/inference_VAR.sh configs/var/varq/base/VAR-VARQ-8.json scripts/output/var

Infinity:

export INFINITY_MODEL_PATH=/path/to/infinity_model
export INFINITY_TEXT_ENCODER_CKPT=/path/to/text_encoder
export INFINITY_PN=1M
bash scripts/inference_Infinity.sh \
  configs/infinity/varq/base/Infinity-VARQ-8.json \
  "a cinematic photograph of a red fox in snow" \
  scripts/output/infinity.png

InfinityStar:

export INFINITYSTAR_CHECKPOINTS_DIR=/path/to/infinitystar/checkpoints
bash scripts/inference_InfinityStar.sh \
  configs/infinitystar/varq/base/InfinityStar-VARQ-8.json \
  --output scripts/output/infinitystar_varq_demo.mp4

Self-Forcing runs its upstream inference.py by default. Arguments after -- are forwarded directly to that script:

bash scripts/inference_SelfForcing.sh \
  configs/self_forcing/varq/base/SF-VARQ-4.json \
  -- <upstream inference arguments>

LongLive follows the same pattern:

bash scripts/inference_LongLive.sh \
  configs/longlive/varq/base/LL-VARQ-4.json \
  -- <upstream inference arguments>

LiveTalk:

bash scripts/inference_LiveTalk.sh 4 -- \
  --checkpoint_root /path/to/livetalk/checkpoints \
  --output scripts/output/livetalk_varq_demo.mp4

If a required third-party checkout is missing, launchers fail early and print the expected third_party/<repo> path.

🧪 Evaluation

VAR evaluation:

bash scripts/eval_VAR.sh \
  configs/var/varq/base/VAR-VARQ-8.json \
  /path/to/VIRTUAL_imagenet256_labeled.npz

Infinity evaluation:

bash scripts/eval_Infinity.sh geneval \
  configs/infinity/varq/base/Infinity-VARQ-8.json

Supported Infinity evaluation tasks are geneval, dpg, and imagereward.

⚙️ Configs

Curated JSON configs live under:

configs/<backend>/<family>/<topic>/*.json

The quantization block is intentionally backend-agnostic and can be reused when calling install_varq_hooks(...) directly. Full JSON files are still kept per backend because model loaders use different qkv_format, sequence layout, image/video schedule, and grouping defaults.

See docs/configs.md for a concise field reference.

Retained public configs include:

For next-frame video backends such as Self-Forcing and LongLive, max_scale_seq_len=1560 is the default grouping unit. If compression_ratio is omitted, it defaults to 1, so the group length is 1560. compression_ratio=3 represents grouping the full 4680-token chunk.

✅ Development Checks

Lightweight tests are public and do not require third-party model repositories, checkpoints, or datasets:

python -m py_compile VAR_Q/*.py VAR_Q/hooks/*.py ablation/*.py scripts/*.py
pytest tests/

The smoke tests cover pack/unpack, quant/dequant shape checks, config loading, and mock runtime hook installation/removal.

🗺️ TODO

• Support more visual autoregressive backends.
• Improve GPU memory fragmentation behavior during long generation.
• Add more backend version signatures for robust hook detection.
• Extend fused kernels and attention-path workspace reuse for additional backends.
• Publish standardized end-to-end memory and throughput benchmarks.
• Add native adapters for additional autoregressive video models.
• Add config inheritance/snippet support so repeated quantization blocks can be shared more compactly.

📚 Citation

If this repository is useful for your research, please cite the VAR-Q paper. The BibTeX entry will be added after the final publication metadata is available.

📄 License

This repository is released under the MIT License. Third-party model repositories, checkpoints, datasets, and generated assets are governed by their own licenses.