File size: 4,743 Bytes
6e668dc | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 112 113 114 115 116 117 | """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 # ms/step
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)
# one new token per request, each token-0 of its own block -> partial block,
# per-token K path (exactly how a decode step grows the trailing block).
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")
|