Speculative Decoding Kv-Cache Optimization

# Speculative Decoding & KV-Cache Optimization

## Introduction & Motivation

Autoregressive language model inference is fundamentally sequential: generating each new token requires a full forward pass through the model conditioned on every token generated so far, and this forward pass cannot begin until the previous token has been sampled, making token-by-token generation latency-bound rather than throughput-bound for a single request. This sequential dependency means that, unlike training (where all positions in a sequence can be processed in parallel because the full target sequence is already known), inference cannot trivially exploit the massive parallelism available on modern accelerators, leaving much of a GPU's compute capacity idle while it waits on memory bandwidth to stream each layer's weights for a single token's forward pass.

Two complementary lines of engineering have emerged to address this bottleneck. Speculative decoding attacks the sequential dependency directly, using a small, fast draft model to propose several tokens ahead, which the large target model then verifies in a single parallel forward pass, converting what would have been several slow sequential steps into one larger, more parallelism-friendly step whenever the draft's guesses are accepted. KV-cache optimization attacks the memory bottleneck: because every previously generated token's key and value projections must be retained and re-read at every subsequent generation step, the key-value (KV) cache grows linearly with sequence length and can, at long context lengths or high batch sizes, dominate both memory footprint and the memory bandwidth consumed per generated token, motivating techniques such as PagedAttention, multi-query and grouped-query attention, and KV-cache quantization.

Together, these techniques are the primary drivers behind the large gap between a language model's theoretical maximum throughput and its naive, unoptimized inference throughput, and are now standard components of essentially every production large language model serving system, directly determining the cost and latency users experience when interacting with deployed language models at scale.

## Core Concepts & Theory

The core insight behind speculative decoding is that verifying a proposed continuation of several tokens is computationally almost as cheap as generating a single token, because both operations require one forward pass through the target model, and a forward pass over multiple positions in parallel is bottlenecked by the same memory bandwidth cost of loading the model's weights, largely independent of how many tokens are processed in that pass (up to the point where compute, rather than memory bandwidth, becomes the limiting factor). Verifying k draft tokens in parallel costs roughly the same wall-clock time as generating a single token normally would, provided the batch of k tokens is still small enough that the operation remains memory-bandwidth-bound.

Given this asymmetry, speculative decoding works by having a small, cheap draft model autoregressively generate a short sequence of candidate tokens, then submitting this entire candidate sequence to the large target model in a single parallel forward pass, which computes the target model's true next-token probability distribution at every position in the candidate sequence simultaneously. An acceptance procedure then compares the draft model's proposed tokens against the target model's actual distribution at each position, accepting tokens that the target model would also have found sufficiently likely and rejecting (and resampling) the first token where the two models diverge meaningfully, guaranteeing that the final output distribution is provably identical to what the target model would have produced through standard autoregressive sampling alone.

KV-cache optimization addresses a separate, memory-centric bottleneck: the attention mechanism's key and value vectors for every previous token must be cached (rather than recomputed) to avoid quadratic recomputation cost during autoregressive generation, but this cache grows linearly with both sequence length and batch size, and every generation step must read the entire cache from GPU memory, meaning that at long context lengths, the time spent reading the KV cache from memory (rather than compute) dominates per-token generation latency, making cache size and memory-access efficiency, rather than raw floating-point throughput, the binding constraint.

## Mathematical Formulation

Speculative decoding's correctness guarantee rests on a rejection-sampling argument. Let q(x) denote the draft model's proposal distribution over the next token and p(x) denote the target model's true distribution. For a token x proposed by the draft model, it is accepted with probability equal to the minimum of one and the ratio of the target model's probability to the draft model's probability for that same token:

$$ P_{\mathrm{accept}}(x) = \min\!\left(1, \; \frac{p(x)}{q(x)} ight) $$

If the proposed token is rejected (which happens with probability 1 minus the acceptance probability), a replacement token is instead sampled from an adjusted residual distribution, constructed by subtracting the draft distribution from the target distribution wherever the target exceeds the draft, and re-normalizing the positive remainder:

$$ p_{\mathrm{residual}}(x) = \frac{\max\!\big(0, \, p(x) - q(x)\big)}{\sum_{x'} \max\!\big(0, \, p(x') - q(x')\big)} $$

This accept-or-resample procedure can be shown to produce samples distributed exactly according to p(x), the target model's true distribution, regardless of how good or poor the draft model q(x) is, which is the key theoretical property that makes speculative decoding a lossless acceleration technique rather than an approximation. The expected number of tokens accepted per speculative round, given an average per-token acceptance rate alpha and a draft length of k tokens proposed per round, follows a simple geometric-style expectation:

$$ \mathbb{E}[ ext{tokens accepted}] = \frac{1 - \alpha^{k+1}}{1 - \alpha} $$

The KV cache's memory footprint for a single sequence, which must be read from memory at every subsequent generation step, scales with the number of transformer layers L, the number of key/value attention heads H_kv, the per-head dimension d_head, the current sequence length T, and 2 bytes per parameter for half-precision storage, doubled again for both keys and values:

$$ ext{KV cache bytes} = 2 imes L imes H_{kv} imes d_{head} imes T imes ext{bytes\_per\_element} $$

## Advanced Theory & Extensions

Tree-based speculative decoding generalizes the basic linear-chain draft-and-verify procedure by having the draft model propose a branching tree of candidate continuations rather than a single linear sequence, allowing multiple plausible continuations to be verified simultaneously in one parallel forward pass through the target model; because the tree's branches share a common prefix, the target model can process the entire tree in a single batched forward pass using an appropriately masked attention pattern, increasing the expected number of accepted tokens per verification round beyond what a single linear draft sequence could achieve, at the cost of a more complex attention mask and verification bookkeeping.

Self-speculative and layer-skipping methods eliminate the need for a wholly separate draft model by instead using a cheaper approximation of the target model itself as the draft, for example by exiting early from a subset of the target model's own layers, or by periodically skipping layers during a fast draft pass and using the full model only for verification; this removes the engineering burden of training, hosting, and keeping a separate draft model synchronized with the target model's vocabulary and tokenizer, at some cost in achievable speedup compared to a purpose-trained, well-matched draft model.

Medusa and related methods replace the separate autoregressive draft model with multiple additional prediction heads attached directly to the target model's final hidden states, each head trained to predict a token several positions ahead in a single forward pass, producing multiple candidate continuations without any additional autoregressive draft-generation steps at all, trading some acceptance-rate quality (since the extra heads are typically less accurate than a true sequential draft model) for a substantially simpler and lower-latency draft-generation step.

On the KV-cache side, multi-query attention (MQA) and grouped-query attention (GQA) reduce cache size by sharing key and value projections across multiple query heads (all query heads in MQA, or groups of query heads in GQA), directly shrinking the H_kv term in the cache-size formula without requiring any change to how many distinct query heads the model uses for its attention computation itself, and have become close to universal in modern large language model architectures specifically because of the KV-cache savings they provide at inference time. PagedAttention, introduced in the vLLM serving system, borrows the concept of virtual memory paging from operating systems, allocating the KV cache in small, fixed-size, non-contiguous blocks rather than one large contiguous buffer per sequence, which eliminates memory fragmentation and enables efficient cache sharing between sequences that share a common prefix (such as multiple samples generated from the same prompt), substantially increasing the number of concurrent sequences a given amount of GPU memory can serve.

## Computational Considerations

The achievable speedup from speculative decoding is fundamentally bounded by the draft model's acceptance rate: a draft model that closely matches the target model's output distribution yields high acceptance rates and large speedups, while a poorly matched draft model wastes computation generating and then discarding proposed tokens that the target model rejects, and in the worst case (a draft model whose proposals are almost always rejected) speculative decoding can actually be slower than standard autoregressive decoding due to the wasted draft-generation overhead, making draft model selection and periodic re-calibration an important practical consideration.

The optimal number of tokens to speculate per round, k, trades off two competing effects: larger k increases the potential number of tokens accepted per verification pass when acceptance rates are high, but also increases the wasted computation when an early token in the draft sequence is rejected (since all subsequently drafted tokens, which depended on that rejected token as context, must be discarded regardless of their own individual plausibility), meaning the optimal k depends heavily on the specific draft-target model pairing's typical acceptance rate and is commonly tuned empirically per deployment.

KV-cache reads dominate inference latency specifically in the memory-bandwidth-bound regime, which is the typical regime for single-sequence or small-batch generation at long context lengths; as batch size increases, the same cache-read cost is amortized across more concurrently generated tokens, shifting the bottleneck from memory bandwidth toward compute, which is precisely why serving systems aggressively batch concurrent requests together (continuous batching) to maximize the useful work extracted from every KV-cache read and every set of model weights loaded from memory.

## Practical Implementation Strategies

Selecting a draft model for speculative decoding typically involves choosing a smaller model from the same architecture family as the target model, ideally one trained on similar data or explicitly distilled from the target model, since draft models that share vocabulary, tokenization, and general output distribution characteristics with the target model achieve substantially higher acceptance rates than an unrelated smaller model; some production systems train a dedicated small draft model specifically to maximize acceptance rate against a particular target model rather than reusing an off-the-shelf smaller model.

Implementing the parallel verification step correctly requires careful attention to numerical precision and sampling temperature consistency between the draft and target model's probability computations, since even small numerical discrepancies between how the two models' logits are computed (for example, differences in mixed-precision rounding) can silently bias the accept/reject procedure away from the theoretically guaranteed exact-match distribution; production implementations typically compute both models' relevant probabilities in the same numerical precision to avoid this issue.

For KV-cache management, adopting a paged or block-based memory allocator (following the PagedAttention pattern) rather than pre-allocating a maximum-length contiguous buffer per sequence is now standard practice in high-throughput serving systems, since pre-allocating for a worst-case maximum sequence length wastes enormous amounts of memory on requests that terminate early, whereas block-based allocation grows the cache incrementally and can immediately reclaim memory from completed or evicted sequences. Prefix caching, in which the KV cache for a shared prompt prefix (such as a common system prompt reused across many requests) is computed once and reused across all requests sharing that prefix, is a further optimization that avoids redundant computation and cache storage for the shared portion of many concurrent requests.

## Benchmark Datasets & Evaluation

Speculative decoding methods are typically evaluated by measuring wall-clock speedup and effective tokens-per-second throughput relative to standard autoregressive decoding, on a fixed target model and hardware configuration, across a range of representative workloads including open-ended chat, code generation, and summarization, since acceptance rates (and therefore achievable speedup) vary substantially by task and output distribution characteristics; code generation, for instance, tends to exhibit unusually high acceptance rates due to its comparatively predictable syntactic structure.

The Spec-Bench benchmark suite standardizes evaluation of speculative decoding methods across multiple task categories (including translation, summarization, question answering, mathematical reasoning, and retrieval-augmented generation) and multiple target model scales, enabling apples-to-apples comparison between different draft model strategies, tree-based verification schemes, and self-speculative methods under consistent hardware and measurement conditions.

KV-cache and serving-system optimizations are commonly evaluated using throughput and latency benchmarks that vary request arrival rate, sequence length distribution, and batch size, reporting metrics such as time-to-first-token (TTFT, dominated by prompt processing), time-per-output-token (TPOT, dominated by the memory-bandwidth cost of reading the growing KV cache), and overall requests-served-per-second under a fixed service-level latency target, with serving systems such as vLLM, TensorRT-LLM, and text-generation-inference (TGI) serving as widely used reference implementations against which new optimizations are benchmarked.

## Key Challenges & Limitations

Speculative decoding's benefit diminishes substantially at large batch sizes, since verifying a batch of draft tokens for many concurrent sequences simultaneously eventually saturates the same compute resources that standard batched autoregressive decoding would also consume, meaning the memory-bandwidth-bound regime in which speculative decoding provides its largest wins (small-batch, latency-sensitive single-sequence generation) is precisely the regime in which many high-throughput serving deployments do not primarily operate, requiring careful workload-specific evaluation of whether speculative decoding is worth its added system complexity for a given deployment's batch-size profile.

Maintaining and synchronizing a separate draft model alongside the target model introduces meaningful operational complexity: the draft model must be kept compatible with the target model's tokenizer and vocabulary, must be re-validated whenever the target model is updated (since acceptance rates can degrade substantially if the draft and target models drift apart in their output distributions), and consumes additional GPU memory and serving infrastructure that must be accounted for in capacity planning.

KV-cache quantization, which reduces cache memory footprint by storing keys and values at lower numerical precision (such as 8-bit or 4-bit representations rather than standard 16-bit), introduces a direct tradeoff between memory savings and generation quality, since attention scores computed from quantized keys and values can accumulate numerical error over long sequences, particularly affecting tasks that depend on precise retrieval of specific information from far back in a long context window, requiring careful per-deployment evaluation of acceptable quality degradation against the memory and throughput gains achieved.

## Hyperparameter Tuning

The draft length k (number of tokens speculatively proposed per verification round) is the primary tunable parameter in speculative decoding, and production systems increasingly use adaptive schemes that adjust k dynamically based on recently observed acceptance rates, proposing longer draft sequences when acceptance rates have recently been high and shortening them when acceptance rates drop, rather than committing to a single fixed k for an entire serving session.

For tree-based speculative decoding, the branching factor and depth of the candidate tree at each verification round represent an additional tuning axis, trading increased verification-pass compute cost (from processing a larger tree in a single batched forward pass) against a higher expected number of accepted tokens per round; excessively wide or deep trees eventually become compute-bound rather than memory-bandwidth-bound, eroding the fundamental efficiency advantage that motivates speculative decoding in the first place.

For KV-cache management, the block size used in paged memory allocation schemes trades off internal fragmentation (larger blocks waste more memory when a sequence's cache doesn't fill an entire block) against allocator bookkeeping overhead (smaller blocks require managing many more allocation units), with typical production block sizes ranging from roughly 16 to 256 tokens per block depending on the specific serving system and expected workload characteristics; the number of GQA key-value head groups similarly trades cache size reduction against a small amount of representational capacity loss relative to full multi-head attention.

## Real-World Applications & Case Studies

Speculative decoding is deployed in essentially all major commercial large language model serving infrastructures as a standard latency-reduction technique for interactive, low-batch-size use cases such as conversational chat interfaces, where reducing the wall-clock time to generate each response directly improves perceived responsiveness; reported speedups in production deployments commonly range from roughly 2x to 3x for well-matched draft-target model pairs on typical conversational workloads.

Code-completion and code-generation products have been particularly strong beneficiaries of speculative decoding, since code's comparatively predictable local syntactic structure (common patterns like closing brackets, repeated variable names, and boilerplate constructs) yields unusually high draft-acceptance rates, and low-latency response is especially valued in interactive coding-assistant use cases where suggestions must appear with minimal perceptible delay as a developer types.

PagedAttention and its associated continuous-batching serving architecture, first popularized by the vLLM project, has been broadly adopted across the open-source and commercial LLM serving ecosystem, and the general pattern of block-based KV-cache management combined with prefix caching for shared system prompts is now considered a baseline requirement for any large language model serving system intended to handle production traffic at meaningful scale and concurrency.

## Integration with Other Methods

Speculative decoding composes naturally with quantization of the target model's own weights (as opposed to KV-cache quantization), since a quantized target model still requires the same verification forward pass, and combining both techniques can compound their respective throughput and memory benefits, though careful evaluation is needed since quantization-induced changes to the target model's output distribution can also shift acceptance rates relative to a draft model tuned against the original full-precision target model.

Retrieval-augmented generation pipelines benefit from prefix caching in particular, since the retrieved context injected into a prompt is often shared across many requests within a short time window (for example, many users asking questions about the same recently retrieved document), allowing the KV cache for that shared retrieved content to be computed once and reused, compounding the throughput benefits of KV-cache optimization with the accuracy benefits of retrieval augmentation.

Mixture-of-Experts target models introduce additional complexity for both speculative decoding and KV-cache management: because only a subset of experts is activated per token, the memory-bandwidth cost structure differs from a dense model, and draft models for MoE targets must be chosen or trained with this different acceptance-rate dynamic in mind, an active area of practical systems research as MoE architectures have become increasingly prevalent in large-scale deployed language models.

## Future Research Directions

An active research direction seeks self-speculative and dynamically adaptive drafting schemes that eliminate the operational burden of maintaining a separate draft model entirely, using techniques such as early-exit prediction heads, layer-skipping, or lightweight n-gram-based lookahead drawn directly from the ongoing generation, aiming to capture much of speculative decoding's speedup benefit without the deployment complexity of a fully separate draft model.

Extending speculative decoding beyond simple next-token verification to speculatively execute entire multi-step reasoning chains or tool-calling sequences in agentic pipelines is an emerging direction, aiming to apply the same draft-then-verify principle at a coarser granularity (verifying whole reasoning steps or tool invocations rather than individual tokens) to accelerate the increasingly common pattern of language models performing extended multi-step reasoning or interacting with external tools before producing a final answer.

On the KV-cache side, research into more aggressive cache compression techniques, including learned cache eviction policies that identify and discard tokens whose cached keys and values contribute little to future attention computations, and hybrid architectures that combine standard attention with constant-memory recurrent or state-space components specifically to bound long-context memory growth, aims to further decouple serving cost from context length as language models are increasingly deployed with context windows spanning hundreds of thousands to millions of tokens.

## Summary & Key Takeaways

Speculative decoding accelerates autoregressive language model inference by having a cheap draft model propose multiple tokens ahead, which the expensive target model then verifies in a single parallel forward pass using a rejection-sampling procedure that provably preserves the target model's exact output distribution; achievable speedup depends fundamentally on the draft model's acceptance rate and is most beneficial in memory-bandwidth-bound, low-batch-size serving regimes such as interactive chat and code completion. KV-cache optimization addresses a complementary bottleneck, using techniques such as multi-query and grouped-query attention, PagedAttention-style block-based memory management, and cache quantization to reduce the memory footprint and bandwidth cost of storing and reading cached keys and values as sequence length and concurrent batch size grow.

Both technique families are now standard components of production LLM serving systems, and their effectiveness is highly workload-dependent, requiring careful tuning of draft length, tree structure, KV-cache block size, and quantization precision against the specific latency, throughput, and quality requirements of a given deployment, with ongoing research aimed at reducing the operational complexity of maintaining separate draft models and further decoupling serving cost from growing context lengths.

Keywords: speculative decoding, draft model verification, rejection sampling acceptance, tree-based speculative decoding, Medusa prediction heads, self-speculative decoding, KV cache optimization, PagedAttention, multi-query attention, grouped-query attention, GQA, MQA, continuous batching, prefix caching, KV cache quantization, memory bandwidth bound inference, time to first token, time per output token, vLLM, LLM inference serving

---

## Appendix: Practical Labs

### Lab 1: Rejection-sampling acceptance rule for speculative decoding

import numpy as np


def rejection_sample_accept(draft_probs, target_probs, draft_token, rng):
    """Implements the core speculative-decoding acceptance rule: a token
    proposed by the draft model is accepted with probability
    min(1, target_prob(token) / draft_prob(token)); if rejected, a
    replacement token is drawn from the normalized residual distribution
    max(0, target_probs - draft_probs)."""
    p_target = target_probs[draft_token]
    p_draft = draft_probs[draft_token]

    accept_prob = min(1.0, p_target / (p_draft + 1e-12))
    accepted = rng.uniform(0, 1) < accept_prob

    if accepted:
        return draft_token, True

    residual = np.maximum(0.0, target_probs - draft_probs)
    residual_sum = residual.sum()
    if residual_sum < 1e-12:
        # Degenerate case: target and draft distributions coincide everywhere
        replacement = int(np.argmax(target_probs))
    else:
        residual = residual / residual_sum
        replacement = int(rng.choice(len(target_probs), p=residual))
    return replacement, False


def test_rejection_sampling_matches_target_distribution():
    rng = np.random.RandomState(0)
    vocab_size = 5

    # A target distribution and a deliberately mismatched draft distribution
    target_probs = np.array([0.5, 0.2, 0.15, 0.1, 0.05])
    draft_probs = np.array([0.1, 0.1, 0.1, 0.1, 0.6])  # draft over-favors token 4

    n_trials = 200_000
    sampled_tokens = np.zeros(n_trials, dtype=int)

    for i in range(n_trials):
        draft_token = rng.choice(vocab_size, p=draft_probs)
        final_token, _ = rejection_sample_accept(draft_probs, target_probs, draft_token, rng)
        sampled_tokens[i] = final_token

    empirical_probs = np.bincount(sampled_tokens, minlength=vocab_size) / n_trials

    print("Target distribution:   ", np.round(target_probs, 3))
    print("Empirical distribution:", np.round(empirical_probs, 3))

    max_abs_error = np.max(np.abs(empirical_probs - target_probs))
    print(f"Max absolute error: {max_abs_error:.4f}")

    assert max_abs_error < 0.01, (
        "Rejection-sampled output distribution should closely match the target "
        "model's true distribution, regardless of the (mismatched) draft distribution"
    )

    print("Speculative decoding rejection-sampling correctness test passed.")


if __name__ == "__main__":
    test_rejection_sampling_matches_target_distribution()

### Lab 2: Expected accepted tokens per round as a function of acceptance rate and draft length

import numpy as np


def expected_accepted_tokens(alpha, k):
    """Closed-form expected number of tokens accepted per speculative round,
    given a per-token acceptance probability alpha and k drafted tokens:
    E[accepted] = (1 - alpha^(k+1)) / (1 - alpha)."""
    if abs(alpha - 1.0) < 1e-9:
        return k + 1.0
    return (1 - alpha ** (k + 1)) / (1 - alpha)


def simulate_accepted_tokens(alpha, k, n_trials, rng):
    """Monte Carlo simulation: at each of k drafted tokens, accept
    independently with probability alpha; the round stops at the first
    rejection. Crucially, the speculative decoding protocol always emits one
    additional token beyond the accepted draft tokens: either a "bonus" token
    sampled from the target model's distribution after all k drafts are
    accepted, or a resampled replacement token drawn from the residual
    distribution at the position of the first rejection. Either way, exactly
    one extra token is always emitted, which is why 1 is added unconditionally
    below rather than only in the all-accepted case."""
    totals = np.zeros(n_trials)
    for i in range(n_trials):
        accepted_count = 0
        for _ in range(k):
            if rng.uniform(0, 1) < alpha:
                accepted_count += 1
            else:
                break
        totals[i] = accepted_count + 1  # bonus token or resampled replacement, always emitted
    return totals.mean()


def test_expected_tokens_formula_matches_simulation_and_scales_with_alpha():
    rng = np.random.RandomState(1)
    k = 5

    alphas = [0.3, 0.6, 0.9]
    closed_form = [expected_accepted_tokens(a, k) for a in alphas]
    simulated = [simulate_accepted_tokens(a, k, 20_000, rng) for a in alphas]

    print(f"{'alpha':>6} | {'closed-form':>12} | {'simulated':>10}")
    for a, cf, sim in zip(alphas, closed_form, simulated):
        print(f"{a:6.2f} | {cf:12.4f} | {sim:10.4f}")
        assert abs(cf - sim) < 0.05, f"Closed-form and simulated expectations should closely agree at alpha={a}"

    # Higher acceptance rate must yield strictly more expected accepted tokens
    assert closed_form[0] < closed_form[1] < closed_form[2]

    # A low acceptance rate should yield close to 1 token per round (mostly rejections)
    assert closed_form[0] < 1.6
    # A high acceptance rate should approach k+1 tokens per round
    assert closed_form[2] > k * 0.7

    print("Expected-accepted-tokens formula validation test passed.")


if __name__ == "__main__":
    test_expected_tokens_formula_matches_simulation_and_scales_with_alpha()

### Lab 3: KV-cache memory footprint under MHA, GQA, and MQA

import numpy as np


def kv_cache_bytes(num_layers, num_kv_heads, head_dim, seq_len, bytes_per_element=2):
    """Computes total KV-cache memory in bytes for a single sequence:
    2 (keys + values) * layers * kv_heads * head_dim * seq_len * bytes_per_element."""
    return 2 * num_layers * num_kv_heads * head_dim * seq_len * bytes_per_element


def test_gqa_and_mqa_reduce_cache_size_proportionally_to_kv_heads():
    num_layers = 32
    num_query_heads = 32
    head_dim = 128
    seq_len = 8192
    bytes_per_element = 2  # fp16

    # Standard multi-head attention: one KV head per query head
    mha_cache = kv_cache_bytes(num_layers, num_query_heads, head_dim, seq_len, bytes_per_element)

    # Grouped-query attention: e.g. 8 KV heads shared across 32 query heads (group size 4)
    num_gqa_groups = 8
    gqa_cache = kv_cache_bytes(num_layers, num_gqa_groups, head_dim, seq_len, bytes_per_element)

    # Multi-query attention: a single shared KV head across all query heads
    mqa_cache = kv_cache_bytes(num_layers, 1, head_dim, seq_len, bytes_per_element)

    mha_gb = mha_cache / (1024 ** 3)
    gqa_gb = gqa_cache / (1024 ** 3)
    mqa_gb = mqa_cache / (1024 ** 3)

    print(f"MHA KV cache: {mha_gb:.3f} GB ({num_query_heads} KV heads)")
    print(f"GQA KV cache: {gqa_gb:.3f} GB ({num_gqa_groups} KV heads)")
    print(f"MQA KV cache: {mqa_gb:.3f} GB (1 KV head)")

    # Cache size should scale exactly linearly with number of KV heads
    assert np.isclose(mha_cache / num_query_heads, gqa_cache / num_gqa_groups)
    assert np.isclose(mha_cache / num_query_heads, mqa_cache / 1)

    # GQA with group size 4 (32/8) should give exactly a 4x reduction vs MHA
    assert np.isclose(mha_cache / gqa_cache, num_query_heads / num_gqa_groups)

    # MQA should give exactly a 32x reduction vs MHA (one head instead of 32)
    assert np.isclose(mha_cache / mqa_cache, num_query_heads)

    assert mqa_cache < gqa_cache < mha_cache, "Cache size should strictly decrease: MHA > GQA > MQA"

    print("KV-cache head-sharing reduction test passed.")


if __name__ == "__main__":
    test_gqa_and_mqa_reduce_cache_size_proportionally_to_kv_heads()

### Lab 4: Simulating paged KV-cache allocation vs. contiguous pre-allocation

import numpy as np


def contiguous_allocation_waste(sequence_lengths, max_seq_len):
    """Simulates naive contiguous allocation: every sequence reserves a
    buffer sized for the maximum possible sequence length up front,
    regardless of how many tokens it actually ends up generating."""
    total_allocated = len(sequence_lengths) * max_seq_len
    total_used = sum(sequence_lengths)
    wasted = total_allocated - total_used
    return total_allocated, total_used, wasted


def paged_allocation_waste(sequence_lengths, block_size):
    """Simulates PagedAttention-style block allocation: each sequence
    allocates only as many fixed-size blocks as needed to hold its actual
    (so-far-generated) tokens, rounded up to the nearest whole block."""
    total_allocated = 0
    total_used = sum(sequence_lengths)
    for length in sequence_lengths:
        num_blocks = -(-length // block_size)  # ceiling division
        total_allocated += num_blocks * block_size
    wasted = total_allocated - total_used
    return total_allocated, total_used, wasted


def test_paged_allocation_wastes_far_less_memory_than_contiguous():
    rng = np.random.RandomState(4)
    n_sequences = 200
    max_seq_len = 4096

    # Realistic scenario: most sequences terminate much earlier than the
    # maximum possible length (a typical heavy-tailed length distribution)
    sequence_lengths = np.clip(
        rng.exponential(scale=300, size=n_sequences).astype(int) + 10,
        10, max_seq_len
    )

    contig_alloc, contig_used, contig_waste = contiguous_allocation_waste(
        sequence_lengths, max_seq_len
    )

    block_size = 16
    paged_alloc, paged_used, paged_waste = paged_allocation_waste(
        sequence_lengths, block_size
    )

    contig_waste_pct = contig_waste / contig_alloc * 100
    paged_waste_pct = paged_waste / paged_alloc * 100

    print(f"Contiguous allocation: {contig_alloc:,} tokens allocated, "
          f"{contig_used:,} used, {contig_waste_pct:.1f}% wasted")
    print(f"Paged allocation (block={block_size}): {paged_alloc:,} tokens allocated, "
          f"{paged_used:,} used, {paged_waste_pct:.1f}% wasted")

    assert contig_used == paged_used, "Both schemes must serve the same actual token usage"
    assert paged_alloc < contig_alloc, "Paged allocation must reserve far less total memory"
    assert paged_waste_pct < contig_waste_pct, "Paged allocation must waste a smaller fraction of allocated memory"

    # Paged waste should be small and bounded by roughly block_size per sequence
    max_possible_paged_waste = n_sequences * block_size
    assert paged_waste <= max_possible_paged_waste

    print(f"Memory savings from paged allocation: {(1 - paged_alloc / contig_alloc) * 100:.1f}%")
    print("Paged vs. contiguous KV-cache allocation test passed.")


if __name__ == "__main__":
    test_paged_allocation_wastes_far_less_memory_than_contiguous()