| """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) |
| v = v_dense[b, :, :L, :].repeat_interleave(group, dim=0) |
| qb = q[b].unsqueeze(1) |
| scores = (qb.float() @ k.float().transpose(-1, -2)) * SM |
| 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 |
|
|
|
|
| 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) |
| |
| 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}") |
|
|
