Reproduce “QUOKA: Query-Oriented KV Selection For Efficient LLM Prefill” on OpenResearch

About this paper

Dalton Jones, Junyoung Park, Matthew Morse, Mingu Lee, Chris Lott, Harper Langston

QUOKA, developed by Qualcomm AI Research, introduces a training-free and hardware-agnostic sparse attention algorithm that accelerates large language model prefill inference, particularly with chunked prefill. The method achieves up to 5x attention speedup and a 3x reduction in time-to-first-token on Nvidia A100 GPUs, while maintaining near-baseline accuracy across various long-context benchmarks.

Key insights

• Queries with low cosine similarity to the mean query vector are critical, as they interact more broadly with keys and contribute most significantly to final attention logits.
• Cosine similarity serves as a more effective, bounded, and geometry-aware proxy for attention weights compared to the traditional dot product, improving subselection quality by over 10%.
• Maximum aggregation across queries is superior to averaging, as it preserves rare but crucial outlier query-key interactions that are essential for maintaining accuracy in sparse attention.

Summary by alphaXiv.