Design efficient Transformer inference with KV cache

Q: Design efficient Transformer inference with KV cache

This question evaluates an engineer's competency in ML systems engineering, specifically understanding of decoder-only Transformer inference, KV cache semantics and which tensors are cached per layer, attention mechanics, memory layout, and correctness under branching and long-context workloads.

Q: How do I approach ML System Design interview questions?

ML System Design questions require understanding of core concepts and practice. PracHub provides solutions with explanations to help you master ml system design interviews.

Question

You are implementing autoregressive inference for a decoder-only Transformer.

Explain what the KV cache is , what tensors are cached per layer, and how it changes computation during incremental decoding.
Describe an implementation plan for KV caching that supports:

Variable sequence lengths in a batch
Beam search or speculative decoding (where sequences can branch)
Long contexts (e.g., 32k–128k tokens)

Discuss key performance considerations:

Memory layout and writes/reads when appending new K/V
Avoiding reallocation/copies
Interaction with fused attention kernels (e.g., FlashAttention-style)
Precision choices (fp16/bf16/int8) for cache

What are common bugs or correctness pitfalls when adding a KV cache (masking, position encodings/RoPE, shape mismatches, etc.)?

Design efficient Transformer inference with KV cache

Quick Overview

Solution

Comments (0)