File size: 3,822 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 | """Microbenchmark: fused paged INT4-KIVI decode vs the dense gather+attend path.
Same packed cache for both; times one decode step (single query token per
request) over a sweep of context lengths and batch sizes. This is the decode-
speed future-work item: the fused path must beat (or at least not regress) the
dense whole-context dequant that materializes (B,H,max_seq,D) bf16 every step.
Run with the vLLM venv from /tmp (avoid vllm package shadowing):
cd /tmp && /home/alex/poolside-hackathon-kv-quant/.venv-vllm/bin/python \
.../scripts/bench_paged_decode.py
"""
from __future__ import annotations
import math
import time
import torch
from vllm.v1.attention.ops.triton_int4_kivi import (
int4_kivi_gather_dequant,
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_cache(B, L, seed=0):
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)
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)
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)
return kv_cache, block_table, seq_lens
def dense_decode(q, kv_cache, block_table, seq_lens):
"""Reproduce the backend's dense path: gather-dequant whole cache + SDPA."""
B = q.shape[0]
max_seq = int(seq_lens.max().item())
k_dense, v_dense = int4_kivi_gather_dequant(
kv_cache, block_table, seq_lens, D, HK, max_seq
)
group = HQ // HK
out = torch.empty(B, HQ, D, dtype=torch.bfloat16, device=DEV)
for b in range(B):
L = int(seq_lens[b].item())
k = k_dense[b, :, :L, :].repeat_interleave(group, dim=0) # [HQ,L,D]
v = v_dense[b, :, :L, :].repeat_interleave(group, dim=0)
qb = q[b].unsqueeze(1) # [HQ,1,D]
scores = (qb.float() @ k.float().transpose(-1, -2)) * SM # [HQ,1,L]
p = torch.softmax(scores, dim=-1)
out[b] = (p @ v.float()).squeeze(1).to(torch.bfloat16)
return out
def timeit(fn, iters=30, warmup=5):
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
if __name__ == "__main__":
print(f"{'B':>3} {'ctx':>7} | {'dense ms':>9} {'fused ms':>9} {'speedup':>8}")
for B, L in [(1, 4096), (1, 12000), (1, 32000),
(8, 4096), (8, 12000), (16, 8000), (32, 4096)]:
kv_cache, bt, sl = build_cache(B, L)
q = torch.randn(B, HQ, D, device=DEV, dtype=torch.bfloat16)
# correctness sanity at this size
ref = dense_decode(q, kv_cache, bt, sl)
fus = int4_kivi_paged_decode(q, kv_cache, bt, sl, SM)
d = (ref.float() - fus.float()).abs().max().item()
dense_ms = timeit(lambda: dense_decode(q, kv_cache, bt, sl))
fused_ms = timeit(lambda: int4_kivi_paged_decode(q, kv_cache, bt, sl, SM))
print(f"{B:>3} {L:>7} | {dense_ms:9.3f} {fused_ms:9.3f} "
f"{dense_ms/fused_ms:7.2f}x max|Δ|={d:.1e}")
|