PROGRESS.md

# 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)
```bash
# 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.
```bash
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)
```bash
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
```bash
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`)
```bash
# 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
```bash
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.