FlashHead
A vLLM plugin that replaces a language model's dense classification head with a two-stage clustered retrieval pipeline, cutting scored vocabulary tokens roughly tenfold for faster, training-free inference.
The dense classification head can account for up to 60% of parameters in small LLMs and roughly half of decode-step compute. FlashHead replaces it with a two-stage retrieval pipeline that reframes head computation as searching clustered token embeddings instead of scoring the full vocabulary at every step, without retraining the model, and integrates through vLLM's official plugin entry point with no source patches or custom Docker image.
Equal-sized clustering groups token embeddings into balanced clusters for predictable memory access, and multi-probe retrieval checks several cluster centroids instead of committing to one, keeping near-perfect recall of the correct next token. It supports both greedy and sampling decoding, choosing clusters proportionally to centroid probabilities to preserve the output distribution, and applies low-precision quantization only to the coarse first stage so the second stage keeps accuracy.
Installed with pip; requires Python 3.10 or later and vLLM 0.14.0 or later.