TL;DR

Reese Chong thêm KV caching và INT8 KV quantization vào transformer inference dựng tay trong mni-ml/framework — một ML library Rust + CUDA viết từ đầu. Kết quả trên prompt 4 token, sinh 252 token: throughput đi từ 0.76 tok/s (baseline) lên 27.29 tok/s (INT8), tức ~35.8× nhanh hơn end-to-end, đồng thời cắt 3.78× dung lượng KV cache mà không mất tốc độ.

What's new

Đây không phải một paper hay feature của vLLM/TensorRT. Đây là một indie hacker bỏ hai kỹ thuật kinh điển của LLM inference (KV cache + INT8 KV quantization) vào một ML framework tự viết bằng Rust với kernel CUDA thủ công, API TypeScript qua N-API. Repo mni-ml/framework hiện 456★. Có cả demo trên trình duyệt để bạn chạy thử.

Why it matters

KV cache và INT8 quantization là tiêu chuẩn trong các production stack như vLLM, TGI, TensorRT-LLM. Nhưng phần lớn developer dùng chúng như black box. Reese làm ngược lại: dựng toàn bộ vertical stack — autograd, tensor, CUDA kernel, attention, cache — từ con số 0. Với ai đang học systems ML hoặc muốn rời PyTorch để hiểu GPU thật sự chạy gì, đây là một trong những reference implementation đầy đủ nhất còn đọc được.

Technical facts

Benchmark trên prompt 4 token, sinh 252 token:

ConfigurationThroughputKV cache memory
Baseline (no cache)0.76 tok/s—
KV cache FP3227.21 tok/s4.50 MB
KV cache INT827.29 tok/s1.19 MB
  • Speedup end-to-end từ KV caching: ~35.8×
  • INT8 giữ nguyên tốc độ (27.29 vs 27.21 tok/s) trong khi giảm bộ nhớ cache 3.78×
  • Stack: Rust backend + CUDA kernel, TypeScript API, hỗ trợ thêm WebGPU
  • Model đi kèm: mni-ml/transformer — 12M params, BPE tokenizer, train trên TinyStories hoặc YouTube-Commons

Cách KV cache tạo ra tốc độ

Không có cache, mỗi bước sinh token mới, model tính lại key và value cho toàn bộ token đã sinh — mặc dù chúng không đổi vì mask ngăn token tương lai. Đó là tính toán thừa O(n²). KV cache lưu lại K và V của token cũ, bước sau chỉ tính thêm 1 K và 1 V mới. Như HuggingFace đo trên Tesla T4 với 1000 token: 11.9s có cache vs 56.2s không cache — khoảng 5× trong điều kiện có cuBLAS tối ưu sẵn. Baseline của Reese chậm hơn nên khoảng cách lên tới ~35×.

Tại sao INT8 mà không mất tốc

INT8 dùng 1 byte/element so với FP32 4 bytes/element — lý thuyết giảm 4× bộ nhớ. Reese đo 3.78× do overhead metadata scale/zero-point. Điều quan trọng là tốc độ không đổi vì bottleneck trong inference là memory bandwidth, không phải compute — nhét cache nhỏ hơn vào GPU memory đồng nghĩa dequantize khi cần vẫn rẻ hơn move data FP32 qua PCIe/HBM. Arxiv 2601.04719 xác nhận INT8 KV cache có error per-element < 0.004, attention-score error < 0.1 cho head 8K — gần như lossless.

Comparison với production stack

  • vLLM: hỗ trợ FP8 KV cache từ 0.3.0, cần driver + hardware recent
  • TensorRT-LLM (NVIDIA): INT8 và FP8 KV cache, chỉ chạy NVIDIA
  • lmdeploy: INT4/INT8 KV cache với group-wise quantization
  • mni-ml/framework: INT8 per-tensor, kernel viết tay, hỗ trợ cả CUDA và WebGPU, MIT-style, dưới 10K LoC đọc được

Production stack có batching, paged attention, speculative decoding — Reese chưa ship những thứ đó. Đổi lại, đây là stack duy nhất bạn có thể đọc hết trong một cuối tuần.

Use cases

  • Học systems ML: đọc từ autograd đến CUDA kernel trong một repo
  • Tham khảo kernel thủ công: CUDA attention không dùng cuBLAS/cuDNN
  • Rust ML ecosystem: xác nhận N-API + Rust backend cho JS tooling là khả thi
  • Inference trên edge: giảm 3.78× KV cache giúp serve ngữ cảnh dài hơn với GPU nhỏ

Limitations & pricing

  • Open-source, không phí. Model 12M params — xa scale frontier, số liệu có thể không extrapolate tuyến tính cho 7B+ LLM nơi attention thực sự memory-bound
  • Chưa có batched inference, paged attention, flash attention, speculative decoding
  • INT8 dùng symmetric per-tensor scaling — có thể drift nhẹ trên context dài, tweet không kèm accuracy benchmark
  • CUDA path chỉ NVIDIA; WebGPU chéo vendor nhưng chưa có benchmark

What's next

Curriculum mni.ml đã có 5 bài: Intro, Gradient Descent, Optimizing the Math, Transformers, LLM Inference. Theo quỹ đạo này, những bài tiếp theo khả năng cao sẽ đi vào batched inference, flash attention, hoặc paged KV cache — đúng hướng production thực sự. Nếu bạn muốn học LLM inference mà không bị lạc trong 100K dòng vLLM, đây là tài liệu đáng bookmark.

Nguồn: tweet gốc của Reese Chong, demo KV cache, mni-ml/framework, HuggingFace KV cache quantization, arxiv 2601.04719.