Vortex: Efficient and Programmable Sparse Attention Serving for AI Agents

Researchers and AI agents often struggle to design and test new ways to make AI models faster by using "sparse attention." Sparse attention aims to reduce the massive computational cost of processing long sequences. It does this by only attending to a subset of key-value (KV) pairs rather than the entire history.

Currently, implementing these new methods is an engineering nightmare. Moving from a theoretical idea to a high-performance kernel requires navigating complex memory layouts. Developers must often write thousands of lines of low-level CUDA code (a language for parallel computing on GPUs). This creates a bottleneck. Both humans and autonomous AI agents are slowed down by the friction of deployment and evaluation.

Vortex changes this by providing a toolkit to write these methods in Python. It then runs them extremely fast on real hardware. The breakthrough is not just the speed. It enables AI agents to actually discover better ways to speed up AI models through autonomous experimentation.

The friction of paged memory

The primary obstacle to efficient sparse attention is how modern LLM serving systems manage memory. To prevent fragmentation (wasted memory gaps), systems like vLLM use paged attention. This stores the KV cache in physically non-contiguous blocks. This is illustrated in, where a logical sequence maps to various disjoint physical pages.

Standard tensor frameworks like PyTorch assume contiguous batch layouts. When you implement a sparse attention algorithm on a paged layout, these memory assumptions break. You cannot simply slice a tensor to get a sequence. Instead, you must deal with indirect addressing and complex pointer arithmetic.

Existing attempts to solve this usually fall into two traps. Some generalize the attention operator to paged layouts. However, they only work for static sparsity patterns where the indices are known upfront. Others provide specialized kernels for specific algorithms but lack programmability. For example, the authors note that implementing one algorithm like Double Sparse in SGLang can require roughly 2,000 lines of code. Much of this is just re-implementing basic math operators over paged tensors. This lack of compositionality means every new idea requires a massive engineering lift.

A dual-stage programmable abstraction

Vortex solves this by decoupling algorithmic logic from physical memory management. It uses two main layers: vFlow and vTensor.

First, the programmer uses vFlow, a Python-embedded frontend. vFlow adopts a single-request mental model. This allows researchers to think in terms of contiguous tensors. It breaks the attention process into two distinct stages: 1. Query-independent preprocessing: This happens during cache construction. The system computes summaries, such as block centroids (representative average points), which are stored in a named cache. These are reused across decoding steps. 2. Query-dependent execution: This is the dynamic part. The system uses the current query to compute relevance scores. It then selects the top-$k$ blocks and performs the actual attention. This logic uses high-level primitives like GeMM (matrix multiplication), TopK, and Softmax, as seen in the block top-k example in .

Second, the vTensor backend interprets these vFlow programs. vTensor extends PyTorch tensors with layout metadata. This includes batch size, pointer arrays, and index structures. This allows the system to maintain a logical view of contiguous tensors. Meanwhile, the underlying engine handles the heavy lifting of gathering data from non-contiguous paged layouts.

To ensure performance, the authors implement several system-level optimizations. They include a workload planner for balanced GPU scheduling and kernel fusion (combining multiple operations into one) to minimize memory traffic. They also use a "stochastic radix top-k" algorithm. This is a lossy variant that uses early termination. It trades a tiny bit of exactness for significant speedups in the indexing stage.

Speedups that scale to 229B parameters

The authors demonstrate that Vortex delivers real-world throughput gains that rival or exceed dense attention.

On the AMC23 benchmark, the paper reports that block top-k attention achieves up to 3.60$\times$ higher throughput than full attention. The Quest algorithm achieves 2.98$\times$ higher throughput. These gains occur without sacrificing accuracy. For user-side latency (the delay experienced by the end user) on long-input workloads, the authors measure P95 latency reductions. These reach up to 11.7$\times$ for block top-k and 12.8$\times$ for Quest on NVIDIA H200 SXM.

The scaling tests provide further evidence. Using NVIDIA B200 GPUs, a rope-aware sparse attention design for the MLA-based GLM-4.7-Flash model achieves a 4.7$\times$ speedup over full attention. They also tested a massive 229B-parameter MiniMax-M2.7 model. This model was served across four GPUs. There, block top-k achieved a 1.37$\times$ throughput improvement.

The paper also highlights the "agentic" capability of the system. In an 18-hour autonomous optimization loop, an AI agent (Claude Code) refined algorithms. It reached up to 3.46$\times$ higher throughput on the AIME24 benchmark while maintaining accuracy, as shown in .

Implementation blind spots

There are clear boundaries to what Vortex currently handles.

First, the system is strictly focused on the decoding phase (the step-by-step generation of tokens). It does not yet support sparse attention optimizations during the prefill phase (the initial processing of the prompt). For workloads where the prompt is exceptionally long, the benefits of Vortex might be delayed.

Second, it lacks training support. Vortex is built for inference serving. It operates on the forward decoding path. It does not handle the backward pass required for training. Researchers cannot yet use it to co-design sparse attention architectures alongside training objectives.

Finally, the abstraction is currently centered on attention-based models with explicit KV caches. The authors admit that extending vTensor to support recurrent or state-space architectures like Mamba is an open challenge. If your production stack moves toward linear-attention models, Vortex cannot yet optimize those specific state-management patterns.

The verdict: A new research instrument

Vortex is a powerful option for anyone working on long-context LLM serving. However, it is currently a tool for optimizing existing architectures. It is not a complete end-to-end training framework.

The real value is the reduction in the "idea-to-production" cycle. By providing a programmable layer that respects paged memory, the authors have changed the nature of the task. Sparse attention is now an algorithmic search task rather than a high-effort kernel engineering task. The fact that AI agents can successfully navigate this design space is significant. It suggests a future where efficient attention mechanisms are discovered by machines.

If you are interested in prototyping, the code is available at https://github.com/Infini-AI-Lab/vortex_torch.

Figures from the paper