INT4-KIVI fused decode — benchmarks & direction

Custom vLLM INT4-KIVI KV-cache decode on Laguna-XS.2 (40 layers, 48 q / 8 kv heads = GQA group 6, head_dim 128, 256-expert MoE / 8 active), B300. Outer worktree branch kv-quant-speedup; vllm submodule branch int4-kivi-speedup (commit b9bf08cb0). Pre-speedup A/B baseline = submodule parent 530247234.

1. Benchmark

Setup (applies to every command)

# Always: vLLM venv, CUDA 12.8, run from /tmp (avoid vllm/ package shadowing).
cd /tmp
ENV="CUDA_HOME=/usr/local/cuda-12.8 VLLM_USE_FLASHINFER_SAMPLER=0"
PY=/home/alex/poolside-hackathon-kv-quant/.venv-vllm/bin/python
S=/home/alex/poolside-hackathon-kv-quant/scripts

Before trusting ANY timing: confirm the GPU is idle — a co-tenant silently corrupts every number (bf16-flash bounces, store shows multi-ms). Correctness is unaffected.

nvidia-smi --query-gpu=memory.used,utilization.gpu --format=csv,noheader
nvidia-smi --query-compute-apps=pid,used_memory --format=csv,noheader

Never run a sweep concurrently with a foreground bench.

Correctness (run first; gate speed work on these)

env $ENV $PY $S/validate_paged_decode.py     # fused decode == dense gather+attend -> ALL PASS
env $ENV $PY $S/validate_store_equiv.py       # store bytes vs golden -> STORE BIT-IDENTICAL
#   (first run with no golden saves it to /tmp/int4_store_golden.pt; --save to refresh)

Kernel microbench

env $ENV $PY $S/bench_paged_decode.py    # fused vs old dense-dequant fallback (speedup x)
env $ENV $PY $S/bench_quant_vs_flash.py  # int4 read/store vs bf16 FlashAttention == THE gap
env $ENV $PY $S/sweep_decode.py          # launch-param sweep (BLOCK_N/warps/stages/waves)

Tuning knobs are env-overridable: VLLM_INT4_DECODE_{BLOCK_N,WARPS,STAGES,WAVES,SPLIT, MAX_SPLIT}. Current defaults BLOCK_N=64, warps=4, stages=3, waves=4.

End-to-end serving (run once per dtype: `KVD=auto` = bf16 ceiling, `KVD=int4_kivi`)

# Batched throughput + executed HumanEval pass@1 (llm.chat, all prompts at once -> B~20):
env $ENV KVD=int4_kivi N=20 PREFIX_TOKENS=12000 MAXNEW=256 $PY $S/longctx_code_serving.py
# Single-stream latency (max_num_seqs=1 -> every decode step is batch-1):
env $ENV KVD=int4_kivi PREFIX_TOKENS=12000 MAXNEW=256 R=8 $PY $S/decode_latency_serving.py

A/B the speedup vs the pre-speedup kernel

cd /home/alex/poolside-hackathon-kv-quant
git -C vllm show 530247234:vllm/v1/attention/ops/triton_int4_kivi.py \
    > /tmp/triton_int4_kivi_OLD.py
cp /tmp/triton_int4_kivi_OLD.py vllm/vllm/v1/attention/ops/triton_int4_kivi.py   # -> OLD
#   ...run any bench above...
git -C vllm checkout -- vllm/v1/attention/ops/triton_int4_kivi.py               # -> restore NEW

Toggle the fused path off (force the dense fallback) with VLLM_INT4_NO_FUSED_DECODE=1.

Current measured numbers (B300, clean GPU, NEW kernel; logs in `results/`)

Correctness: validate_paged_decode ALL PASS (fused==dense, max|Δ|≈2e-3, bit-exact L=1); validate_store_equiv STORE BIT-IDENTICAL.

Fused vs old dense-dequant fallback (bench_paged_decode.py): 9.4–11.9× across B=1..32, ctx 4k–32k.

Decode-attention read vs bf16 FlashAttention (bench_quant_vs_flash.py, one step):

B	ctx	bf16 FA ms	int4 read ms	read×	(pre-speedup)
1	4k	0.027	0.050	1.9×	6.8×
1	12k	0.027	0.115	4.3×	17.8×
1	32k	0.035	0.279	8.0×	35.3×
8	12k	0.068	1.038	15.2×	17.7×
32	4k	0.091	1.262	13.9×	16.4×
Per-token store: 0.024 ms (B=1) → 0.069 ms (B=32), sync-free.

Batched HumanEval serving (longctx_code_serving.py, N=20, 12k long regime):

KV	short pass@1 / gen	long pass@1 / gen
bf16 (auto)	20/20 · 29s	20/20 · 22s
int4 NEW	19/20 · 35s	18/20 · 53s
int4 OLD	19/20 · 31s	16/20 · 53s

Single-stream latency (decode_latency_serving.py, max_num_seqs=1, 12k, 256 tok, R=8 median):

KV	ms/tok	tok/s	vs bf16
bf16 (auto)	22.5	44.4	—
int4 NEW	37.4	26.8	1.66×
int4 OLD	39.9	25.1	1.77×

2. Direction of improvement

State of the gap, by regime — the numbers above locate where work remains; this section names the bottlenecks, it does not propose solutions.

Kernel, batch-1: closed. Decode-attention read is 1.9–8× of bf16 FlashAttention (was 6.8–35×). This is the regime the bf16-tensor-core change targeted.
Kernel, large batch (B≥8): open, dominant. The int4 read is still ~14–17× of bf16 FA and barely moved from the pre-speedup kernel. This is the largest remaining kernel gap. Two structural causes are visible in the geometry: GQA group = 6 forces the QK/PV tl.dot to pad M to 16 (10/16 tensor-core rows unused), and the int4→bf16 unpack is ALU-bound rather than bandwidth-bound (it moves 3.2× fewer bytes yet runs slower).
End-to-end, single-stream latency: improved, not eliminated. int4 decode 39.9 → 37.4 ms/tok (1.07× NEW vs OLD); gap to bf16 1.77× → 1.66×. The int4 e2e penalty (+14.8 ms/tok over bf16) is far larger than the decode-attention compute difference alone — i.e. most of the single-stream e2e gap is per-step int4-backend overhead, not the attention math.
End-to-end, batched throughput: unchanged. int4 NEW ≈ OLD (~53 s). At B≈20 the read sits in the still-slow large-batch regime, and on this 40-layer MoE decode-attention is a small fraction of each step (expert FFN + projections + sampling, all shared/unquantized, dominate), so attention-only changes cannot move batched wall-clock much. The win at this scale is the 3.2× KV-cache memory (longer context / more concurrent streams), not latency.
CUDA graphs: unblocked, not enabled. The decode path is now sync-free, but the metadata builder is still AttentionCGSupport.NEVER; graph capture has not been turned on or measured.