$ show microsoft/vattention
This repository is a research prototype for efficient LLM inference, specifically optimizing attention memory management via the vAttention technique (ASPLOS 2025). It contains three primary components: (1) a modified NVIDIA GPU kernel driver (nvidia-vattn-uvm-driver) that enables sub-2MB page sizes for Unified Virtual Memory (UVM), (2) a fork of the Sarathi-Serve LLM serving framework (sarathi-lean) integrated with vAttention, and (3) custom attention kernel implementations (pod_attn, vattention) using CUTLASS and FlashAttention backends. The system supports multiple attention backends (FlashAttention, FlashInfer, paged, variable-length) and includes extensive benchmarking and plotting scripts for reproducibility of research results. Architecturally, the user-space inference stack uses a scheduler (Orca or Sarathi), a cache engine (vLLM-style or vAttention), and model executors with tensor/pipeline parallelism. The vAttention cache engine replaces traditional paging with CUDA virtual memory allocation, decoupling virtual and physical memory for KV-cache. The custom NVIDIA UVM driver patch is required to support fine-grained page sizes. The repository also includes microbenchmarks for attention kernel performance and page size effects, as well as artifact scripts to reproduce figures from the associated paper. While the codebase demonstrates solid engineering for a research project (clean abstractions, registry pattern, modular design), it is explicitly not production-ready. Key limitations include strict version dependencies (PyTorch 2.3.0, CUDA 12.1), reliance on a modified kernel driver, and incomplete feature parity with production systems like vLLM. The inclusion of a fully vendored CUTLASS library and auto-generated NVIDIA driver headers adds significant size but provides reusable GPU primitives. The research scripts are functional but contain hardcoded paths and are tightly coupled to the experimental setup.
The LLM identified 3 files worth highlighting for reuse:
Implements a memory manager for KV-cache using CUDA virtual memory (CUDA VA) API. Provides functions to init, reserve, step, free, and allocate physical pages. Uses pybind11 to expose a Python interface. Self-contained and could be extracted into a standalone library for dynamic memory allocation in GPU compute applications, particularly for long-context inference.
Provides CpuTimer and CudaTimer context managers for timing operations with CPU/GPU events. Appends results to a MetricStore that aggregates and summarizes timing data. Reusable as a generic profiling utility for any CUDA-based Python application.
Abstract base class for attention backends using a singleton pattern per model layer. Registers concrete implementations (FlashAttention, FlashInfer, vAttention) and manages per-layer timers. Demonstrates a clean extensible design for swapping attention kernels in LLM inference systems.