KVCache-DSL: An MLIR-based Domain-Specific Language for KV Cache Optimization in Large Language Models

KVCache-DSL is an MLIR-based domain-specific language project aimed at addressing key performance issues in KV cache memory management during large language model (LLM) inference. By jointly analyzing and transforming the memory layout, access patterns, and vectorization of KV caches, this project provides an innovative solution for LLM inference optimization.

During the autoregressive generation process of LLMs, KV caches need to store Key and Value tensors for each layer to avoid redundant computations, but this also brings three major pain points:

1. Huge memory footprint: In long-sequence and batch inference scenarios, KV caches can occupy tens or even hundreds of gigabytes of GPU memory;
2. Complex access patterns: Different model architectures (e.g., Transformer, Mamba, RWKV) have significantly different access patterns for KV caches;
3. Coupling between layout and vectorization: Memory layout decisions directly affect SIMD vectorization efficiency, but the two are often optimized separately. Traditional methods treat these three as independent problems, making it difficult to achieve global optimality.

The core design of KVCache-DSL revolves around the joint analysis and transformation methodology, covering three dimensions:

1. Memory Layout

Describe the physical storage structure of KV caches (continuous, paged, custom layouts, etc.) in a declarative manner, making layout decisions first-class citizens that can be analyzed and transformed;

2. Access Patterns

Capture the read/write patterns of KV caches (e.g., query-key matching in attention computation, autoregressive incremental updates, multi-turn dialogue history reuse) via MLIR dialects, supporting targeted optimizations like prefetching and cache alignment;

3. Vectorization

Deeply couple vectorization strategies with memory layout—developers can specify vector width, alignment requirements, etc., and the compiler generates optimal code based on the SIMD features of the target hardware, avoiding performance losses caused by disjoint layout and vectorization.

Choosing MLIR as the infrastructure brings multiple advantages:

• Progressive lowering: Gradually lower from high-level DSL to LLVM IR, with specific analysis and transformation passes insertable at each layer to form a complete optimization pipeline;
• Multi-target support: The unified intermediate representation allows the same DSL to generate code for multiple backends such as CPU, GPU, NPU, without rewriting front-end logic;
• Ecosystem integration: Seamlessly integrate with the existing MLIR ecosystem (e.g., polyhedral optimization for Affine dialect, CUDA/ROCm code generation for GPU dialect).

KVCache-DSL has broad application prospects:

• Inference engine development: Frameworks like vLLM and TensorRT-LLM can integrate the DSL to achieve more flexible KV cache management;
• Model architecture innovation: New attention mechanisms (e.g., linear attention, state space models) can quickly validate KV cache optimization schemes;
• Hardware co-design: Chip manufacturers can define hardware primitives based on the DSL to achieve hardware-software co-optimization.

Technical challenges and future directions for the project:

1. Automatic scheduling: Need stronger AutoTuning support to automatically derive optimal memory layout and access scheduling strategies from high-level DSL;
2. Dynamic shape handling: LLM inference sequence lengths are dynamic, so the DSL needs better support for compile-time optimization of dynamic shapes;
3. Framework integration: Need to solve the graph capture and code generation problems when embedding the DSL into mainstream frameworks like PyTorch and JAX.

KVCache-DSL represents an important direction in the field of LLM inference optimization: by combining compiler technology with domain-specific languages, it transforms KV cache management—originally dependent on manual tuning—into a systematic and reusable engineering practice. As the project evolves, it is expected to become a key component of the next-generation efficient LLM inference infrastructure.