File size: 3,493 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 | """Sweep fused INT4-KIVI decode launch params to pick the best B300 config.
Times only the fused read (int4_kivi_paged_decode) across a small grid of
(BLOCK_N, num_warps, num_stages, split-waves) for a few (B, ctx) shapes and
prints the best per shape. Mutates the module-level tuning globals between
runs (the launcher reads them per call). Run from /tmp with the vLLM venv:
cd /tmp && CUDA_HOME=/usr/local/cuda-12.8 .venv-vllm/bin/python \
.../scripts/sweep_decode.py
"""
from __future__ import annotations
import math
import time
import torch
import vllm.v1.attention.ops.triton_int4_kivi as K
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_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 timeit(fn, iters=40, 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
SHAPES = [(1, 4096), (1, 12000), (1, 32000), (8, 12000), (16, 8000), (32, 4096)]
BLOCK_NS = [64, 128]
WARPS = [2, 4]
STAGES = [2, 3]
WAVES = [1, 2, 4]
if __name__ == "__main__":
for B, L in SHAPES:
kv_cache, bt, sl = build_cache(B, L)
q = torch.randn(B, HQ, D, device=DEV, dtype=torch.bfloat16)
best = (1e9, None)
results = []
for bn in BLOCK_NS:
for w in WARPS:
for st in STAGES:
for wv in WAVES:
K._DECODE_BLOCK_N = bn
K._DECODE_NUM_WARPS = w
K._DECODE_NUM_STAGES = st
K._DECODE_WAVES = wv
try:
ms = timeit(
lambda: int4_kivi_paged_decode(
q, kv_cache, bt, sl, SM
)
)
except Exception as e: # noqa: BLE001
ms = float("nan")
cfg = (bn, w, st, wv)
results.append((ms, cfg))
if ms < best[0]:
best = (ms, cfg)
results.sort()
print(f"\n=== B={B} ctx={L} === best {best[0]:.3f}ms cfg(BLOCK_N,warps,stages,waves)={best[1]}")
for ms, cfg in results[:5]:
print(f" {ms:7.3f}ms {cfg}")
|