eyestech-mla

eyestech-mla is a production-grade, standalone PyTorch reference kernel and memory scaling benchmark suite for DeepSeek Multi-Head Latent Attention (MLA). It demonstrates Query Absorption, Decoupled Rotary Position Embedding (RoPE), and Zero-Decompression KV Cache Streaming for high-throughput LLM inference up to 128k context windows.

📖 Canonical Systems Audit & Mathematical Proof:
For the complete mathematical proof, memory bandwidth derivations, and production serving economics (vLLM / SGLang), read the flagship investigation:
👉 DeepSeek MLA Architecture: How It Cuts KV Cache by 93% published by EyesTech Systems Lab.

⚡ Key Architectural Highlights

In autoregressive large language model serving, the primary hardware bottleneck at 32k–128k sequence lengths is High Bandwidth Memory (HBM) exhaustion, not raw compute FLOPs.

How DeepSeek MLA Achieves 92.97% Memory Compression:

1. Low-Rank Latent Compression ($c_t^{KV}$): Projects multi-head Key and Value tensors into a compact 512-dimensional shared latent subspace ($d_c = 512$).
2. Decoupled Rotary Position Embedding ($k_t^R$): Employs a dedicated 64-dimensional uncompressed RoPE key vector ($d_h^R = 64$), preserving exact positional distance semantics without inflating latent memory ($512 + 64 = 576$ scalars per token).
3. Inference Query Absorption ($q_{\text{absorbed}} = q \cdot W_{UK}$): Leverages associative matrix multiplication during autoregressive decoding. The up-projection matrices are pre-folded into the active Query tensors, allowing attention dot-products directly against compressed latent states with zero decompression overhead.

🚀 Installation

Install via pip:

pip install eyestech-mla

Or install from source:

git clone https://github.com/abhishek2512mishra/deepseek-mla-kvcache.git
cd deepseek-mla-kvcache
pip install .

💻 Quickstart & Code Examples

1. Autoregressive MLA Decoding with Query Absorption

import torch
from eyestech_mla import MultiHeadLatentAttentionDecode

device = "cuda" if torch.cuda.is_available() else "cpu"

# Instantiate DeepSeek MLA kernel matching V2/V3 configurations
mla = MultiHeadLatentAttentionDecode(
    d_model=5120,
    n_heads=128,
    d_head=128,
    d_latent=512,
    d_rope=64
).to(device)

# Simulate 128k context stream at position 1,024
batch_size = 2
seq_len = 1024
h_t = torch.randn(batch_size, 1, 5120, device=device)
cache_latent = torch.randn(batch_size, seq_len, 512, device=device)
cache_rope = torch.randn(batch_size, seq_len, 64, device=device)

# Execute single-token decode step
output, new_latent, new_rope = mla.forward_decode(
    h_t,
    current_pos=seq_len,
    kv_cache_latent=cache_latent,
    kv_cache_rope=cache_rope
)

print(f"Token Output Tensor: {output.shape}")          # [2, 1, 5120]
print(f"Updated Latent Cache: {new_latent.shape}")      # [2, 1025, 512]
print(f"Updated RoPE Cache:   {new_rope.shape}")        # [2, 1025, 64]
print(f"Stored Scalars/Token: {new_latent.shape[-1] + new_rope.shape[-1]}")  # 576

2. KV Cache Scaling Benchmark (CLI)

Run the automated scaling benchmark directly from your terminal:

eyestech-mla-benchmark

3. Programmatic Memory Benchmark API

from eyestech_mla import compute_kv_cache_bytes, format_bytes

# Compute KV cache for 128k context, 60 layers, batch size 8 in FP16
bytes_mla = compute_kv_cache_bytes(
    n_layers=60,
    n_kv_heads=128,
    d_head=128,
    seq_len=131072,
    batch_size=8,
    bytes_per_elem=2,
    is_mla=True,
    d_latent=512,
    d_rope=64
)

bytes_mha = compute_kv_cache_bytes(
    n_layers=60,
    n_kv_heads=128,
    d_head=128,
    seq_len=131072,
    batch_size=8,
    bytes_per_elem=2,
    is_mla=False
)

print(f"DeepSeek MLA Footprint: {format_bytes(bytes_mla)}")  # 67.50 GB
print(f"Frontier MHA Footprint: {format_bytes(bytes_mha)}")  # 3840.00 GB

📊 Benchmark Telemetry (60 Layers, Batch Size = 8, FP16)

========================================================================================
Architecture               | Scalars/Tok | 4k (4096)  | 16k (16384) | 32k (32768) | 64k (65536) | 128k (131072)
----------------------------------------------------------------------------------------
Standard MHA (32-head)     | 8192        | 30.00 GB   | 120.00 GB   | 240.00 GB   | 480.00 GB   | 960.00 GB   
Frontier MHA (128-head)    | 32768       | 120.00 GB  | 480.00 GB   | 960.00 GB   | 1920.00 GB  | 3840.00 GB  
Llama-3 GQA (8-head)       | 2048        | 7.50 GB    | 30.00 GB    | 60.00 GB    | 120.00 GB   | 240.00 GB   
DeepSeek MLA (Absorbed)    | 576         | 2.11 GB    | 8.44 GB     | 16.88 GB    | 33.75 GB    | 67.50 GB    
========================================================================================

📚 Citation & Attribution

If you use this benchmark harness, reference implementation, or mathematical formalization in academic research, benchmark audits, or engineering systems, please cite:

@misc{fischer2026deepseekmla,
  author = {Fischer, Klaus and Ranganathan, Devika},
  title = {DeepSeek MLA Architecture: How It Cuts KV Cache by 93%},
  howpublished = {\url{https://eyestech.in/deepseek-mla-architecture-kv-cache-math/}},
  journal = {EyesTech Systems Research},
  year = {2026},
  note = {EyesTech Systems Lab Hardware Audit Series}
}

⚖️ License

Distributed under the MIT License. Maintained by EyesTech Systems Lab.