File size: 4,847 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 118 119 120 121 | """Validate the fused paged INT4-KIVI decode kernel against the dense path.
Both sides read the SAME packed int4 cache (built with ``int4_kivi_store`` exactly
as vLLM's prefill does), so this isolates KERNEL correctness from quant error:
the fused ``int4_kivi_paged_decode`` must match dequant-the-whole-cache
(``int4_kivi_gather_dequant``) + GQA softmax attention, to fp32 tolerance.
Run with the vLLM venv from a NON-vllm cwd:
cd /tmp && /home/alex/poolside-hackathon-kv-quant/.venv-vllm/bin/python \
/home/alex/poolside-hackathon-kv-quant/.claude/worktrees/kv-quant-decode-speed/scripts/validate_paged_decode.py
"""
from __future__ import annotations
import math
import torch
from vllm.v1.attention.ops.triton_int4_kivi import (
BLOCK,
int4_kivi_gather_dequant,
int4_kivi_paged_decode,
int4_kivi_store,
)
DEV = "cuda"
HQ, HK, D = 48, 8, 128 # Laguna-XS.2 geometry (GQA group 6)
PAGE = 16 # paged block_size (tokens per page)
FULL_DIM = D // 2 + D // 16 # 64 + 8 = 72
def build_cache(seq_lens, seed=0):
"""Build a paged int4 cache + block_table for the given per-request seq_lens,
storing each request's whole sequence prefill-style (monotone slots from a
block boundary -> full blocks become per-channel K, trailing block per-token).
Returns (kv_cache, block_table, seq_lens_t, k_ref, v_ref) where k_ref/v_ref
are the original bf16 K/V (only for sanity, not used as the reference)."""
g = torch.Generator(device=DEV).manual_seed(seed)
B = len(seq_lens)
nblk_per = [(L + PAGE - 1) // PAGE for L in seq_lens]
max_blocks = max(nblk_per)
num_blocks = sum(nblk_per) + 4
kv_cache = torch.zeros(
(num_blocks, 2, PAGE, HK, FULL_DIM), dtype=torch.uint8, device=DEV
)
block_table = torch.zeros((B, max_blocks), dtype=torch.int32, device=DEV)
cursor = 1 # leave block 0 unused to catch base-offset bugs
for b, L in enumerate(seq_lens):
nb = nblk_per[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)
# slot_mapping: token t -> phys_block*PAGE + (t % PAGE)
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_t = torch.tensor(seq_lens, dtype=torch.int32, device=DEV)
return kv_cache, block_table, seq_lens_t
def ref_attend(q, kv_cache, block_table, seq_lens):
"""Dense reference: gather-dequant the whole cache, GQA softmax attention."""
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
) # [B, HK, max_seq, D] bf16
group = HQ // HK
sm = 1.0 / math.sqrt(D)
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, :].float().repeat_interleave(group, dim=0) # [HQ,L,D]
v = v_dense[b, :, :L, :].float().repeat_interleave(group, dim=0)
qb = q[b].float().unsqueeze(1) # [HQ,1,D]
scores = (qb @ k.transpose(-1, -2)) * sm # [HQ,1,L]
p = torch.softmax(scores, dim=-1)
out[b] = (p @ v).squeeze(1).to(torch.bfloat16)
return out
def run_case(seq_lens, seed=0):
kv_cache, bt, sl = build_cache(seq_lens, seed=seed)
g = torch.Generator(device=DEV).manual_seed(seed + 999)
B = len(seq_lens)
q = torch.randn(B, HQ, D, generator=g, device=DEV, dtype=torch.bfloat16)
sm = 1.0 / math.sqrt(D)
ref = ref_attend(q, kv_cache, bt, sl)
fused = int4_kivi_paged_decode(q, kv_cache, bt, sl, sm)
diff = (fused.float() - ref.float()).abs()
rel = diff / (ref.float().abs() + 1e-3)
return diff.max().item(), diff.mean().item(), rel.max().item()
if __name__ == "__main__":
torch.manual_seed(0)
cases = {
"exact-block (L=512)": [512],
"partial-tail (L=500)": [500],
"short (L=33)": [33],
"tiny (L=1)": [1],
"mixed batch": [128, 257, 64, 1000, 16, 999],
"long (L=12000)": [12000],
"long mixed": [12000, 8001, 16000, 4096],
}
ok = True
for name, sl in cases.items():
amax, amean, rmax = run_case(sl, seed=hash(name) % 10000)
tol = 5e-2
passed = amax < tol
ok = ok and passed
flag = "ok " if passed else "FAIL"
print(f"[{flag}] {name:24s} max|Δ|={amax:.4e} mean|Δ|={amean:.2e} relmax={rmax:.2e}")
print("ALL PASS" if ok else "SOME FAILED")
|