MoSKA: Mixture of Shared KV Attention for Efficient Long-Sequence LLM Inference

2025-11-08

IEEE computer architecture letters

AI Summary

To address memory bottlenecks and low GPU utilization caused by KV caching in long-context LLM inference, this paper proposes Mixture of Shared KV Attention (MoSKA). Methodologically, MoSKA explicitly distinguishes request-specific from shared context across concurrent requests to enable cross-request KV cache reuse; reformulates shared-KV attention computation—from memory-bound GEMV operations into compute-bound batched GEMM; and integrates an MoE-inspired sparse attention mechanism with a software-hardware co-designed disaggregated architecture. Experimental results demonstrate that, under high-context-sharing workloads, MoSKA achieves up to 538.7× higher throughput than baseline systems, significantly improving scalability and hardware efficiency for long-sequence inference.

Technology Category

• Machine Learning: Mixture of Experts (MoE)
• Data Mining & Knowledge Management: Scalability, Parallel & Distributed Systems
• Computer Vision: Large Vision Models

Application Category

• Semantics and Knowledge: Methods to enhance, augment, integrate or synergize semantic models such as knowledge graphs and LLMs
• Search and Retrieval-Augmented AI: Large language models for search
• Economics, Online Markets and Human Computation: Cost models of using LLMs in production systems

Abstract

The escalating context length in Large Language Models (LLMs) creates a severe performance bottleneck around the Key-Value (KV) cache, whose memory-bound nature leads to significant GPU under-utilization. This paper introduces Mixture of Shared KV Attention (MoSKA), an architecture that addresses this challenge by exploiting the heterogeneity of context data. It differentiates between per-request unique and massively reused shared sequences. The core of MoSKA is a novel Shared KV Attention mechanism that transforms the attention on shared data from a series of memory-bound GEMV operations into a single, compute-bound GEMM by batching concurrent requests. This is supported by an MoE-inspired sparse attention strategy that prunes the search space and a tailored Disaggregated Infrastructure that specializes hardware for unique and shared data. This comprehensive approach demonstrates a throughput increase of up to 538.7x over baselines in workloads with high context sharing, offering a clear architectural path toward scalable LLM inference.

Problem

• Addresses KV cache bottleneck in long-sequence LLM inference
• Exploits context heterogeneity between unique and shared sequences
• Transforms memory-bound attention into compute-bound operations

Innovation

• Shared KV Attention converts GEMV to GEMM
• MoE-inspired sparse attention prunes search space
• Disaggregated Infrastructure specializes hardware for data