| """How much does the INT4-KIVI quantization pipeline cost vs *only* flash decode? |
| |
| Baseline ("only flash decode"): a bf16 KV cache attended with FlashAttention |
| (flash_attn_varlen_func) — exactly what vLLM's KVD=auto path runs, no quant. |
| Quantized path: our fused INT4-KIVI decode (int4_kivi_paged_decode) reading the |
| packed 4-bit cache, dequant-in-kernel. |
| |
| We isolate the per-decode-step cost on identical shapes: |
| * read/attend: fused int4 decode vs bf16 flash decode |
| * store: per-token int4 quant (int4_kivi_store of 1 new token) — the other |
| half of the quant pipeline that bf16 doesn't pay. |
| |
| Run from /tmp with the vLLM venv: |
| cd /tmp && CUDA_HOME=/usr/local/cuda-12.8 .venv-vllm/bin/python \ |
| .../scripts/bench_quant_vs_flash.py |
| """ |
|
|
| from __future__ import annotations |
|
|
| import math |
| import time |
|
|
| import torch |
|
|
| from vllm.v1.attention.backends.fa_utils import flash_attn_varlen_func |
| from vllm.v1.attention.ops.triton_int4_kivi import ( |
| int4_kivi_paged_decode, |
| int4_kivi_store, |
| ) |
|
|
| DEV = "cuda" |
| HQ, HK, D = 48, 8, 128 |
| PAGE = 16 |
| FULL_DIM = D // 2 + D // 16 |
| SM = 1.0 / math.sqrt(D) |
|
|
|
|
| def build(B, L, seed=0): |
| """Return bf16 dense K/V (varlen-packed) + the equivalent packed int4 cache.""" |
| g = torch.Generator(device=DEV).manual_seed(seed) |
| nb = (L + PAGE - 1) // PAGE |
| num_blocks = B * nb + 4 |
| kv_cache = torch.zeros( |
| (num_blocks, 2, PAGE, HK, FULL_DIM), dtype=torch.uint8, device=DEV |
| ) |
| block_table = torch.zeros((B, nb), dtype=torch.int32, device=DEV) |
| k_pack = torch.empty(B * L, HK, D, dtype=torch.bfloat16, device=DEV) |
| v_pack = torch.empty(B * L, HK, D, dtype=torch.bfloat16, device=DEV) |
| cursor = 1 |
| for b in range(B): |
| phys = list(range(cursor, cursor + nb)) |
| cursor += nb |
| for j, p in enumerate(phys): |
| block_table[b, j] = p |
| k = torch.randn(L, HK, D, generator=g, device=DEV, dtype=torch.bfloat16) |
| v = torch.randn(L, HK, D, generator=g, device=DEV, dtype=torch.bfloat16) |
| k_pack[b * L : (b + 1) * L] = k |
| v_pack[b * L : (b + 1) * L] = v |
| slots = torch.tensor( |
| [phys[t // PAGE] * PAGE + (t % PAGE) for t in range(L)], |
| dtype=torch.int64, device=DEV, |
| ) |
| int4_kivi_store(k, v, kv_cache, slots, D) |
| seq_lens = torch.full((B,), L, dtype=torch.int32, device=DEV) |
| cu_q = torch.arange(B + 1, dtype=torch.int32, device=DEV) |
| cu_k = torch.arange(0, B * L + 1, L, dtype=torch.int32, device=DEV) |
| return kv_cache, block_table, seq_lens, k_pack, v_pack, cu_q, cu_k |
|
|
|
|
| def bf16_flash(q, k_pack, v_pack, cu_q, cu_k, L): |
| return flash_attn_varlen_func( |
| q=q, k=k_pack, v=v_pack, |
| cu_seqlens_q=cu_q, cu_seqlens_k=cu_k, |
| max_seqlen_q=1, max_seqlen_k=L, |
| softmax_scale=SM, causal=True, fa_version=4, |
| ) |
|
|
|
|
| def timeit(fn, iters=50, warmup=10): |
| for _ in range(warmup): |
| fn() |
| torch.cuda.synchronize() |
| t0 = time.perf_counter() |
| for _ in range(iters): |
| fn() |
| torch.cuda.synchronize() |
| return (time.perf_counter() - t0) / iters * 1e3 |
|
|
|
|
| def store_one(kv_cache, B): |
| """Cost of quantizing+storing one new decode token per request.""" |
| g = torch.Generator(device=DEV).manual_seed(7) |
| k = torch.randn(B, HK, D, generator=g, device=DEV, dtype=torch.bfloat16) |
| v = torch.randn(B, HK, D, generator=g, device=DEV, dtype=torch.bfloat16) |
| |
| |
| slots = (torch.arange(B, dtype=torch.int64, device=DEV)) * PAGE |
| return lambda: int4_kivi_store(k, v, kv_cache, slots, D) |
|
|
|
|
| if __name__ == "__main__": |
| print("Per decode step, B300. read = attention over the cached context;") |
| print("store = quantize+write the 1 new token (int4 only). 'pipeline' = read+store.\n") |
| print(f"{'B':>3} {'ctx':>7} | {'bf16 flash':>10} {'int4 read':>10} {'int4 store':>10} " |
| f"{'int4 pipe':>10} | {'read x':>7} {'pipe x':>7}") |
| for B, L in [(1, 4096), (1, 12000), (1, 32000), |
| (8, 4096), (8, 12000), (16, 8000), (32, 4096)]: |
| kv_cache, bt, sl, kp, vp, cu_q, cu_k = build(B, L) |
| q = torch.randn(B, HQ, D, device=DEV, dtype=torch.bfloat16) |
| qf = q.reshape(B, HQ, D) |
| bf16_ms = timeit(lambda: bf16_flash(qf, kp, vp, cu_q, cu_k, L)) |
| int4_read_ms = timeit(lambda: int4_kivi_paged_decode(q, kv_cache, bt, sl, SM)) |
| int4_store_ms = timeit(store_one(kv_cache, B)) |
| pipe = int4_read_ms + int4_store_ms |
| print(f"{B:>3} {L:>7} | {bf16_ms:10.3f} {int4_read_ms:10.3f} {int4_store_ms:10.3f} " |
| f"{pipe:10.3f} | {int4_read_ms/bf16_ms:6.2f}x {pipe/bf16_ms:6.2f}x") |
|
|
