This document is relevant for: Trn2, Trn3

KV Cache Analysis Design#

Overview#

The KV Cache Analysis module compares KV caches between an expected reference model and a target model (vLLM Neuron) to debug accuracy issues at the attention layer level. It optionally includes a baseline (ground truth) model for three-way comparison, enabling separation of dtype-only errors from target-specific errors.

Naming convention (matches logit_validation):

expected: reference to compare against (required)
actual: target under test (vLLM Neuron)
baseline: ground truth for three-way BC analysis (optional); when provided, all errors are relative to baseline since both expected and actual are teacher-forced against it

Integration with Logit Validation: KV cache extraction is integrated into the logit_validation API via the kv_extract_fn(seq_len) parameter. This ensures KV caches are captured during teacher-forced generation, where both models process identical token sequences. The seq_len argument tells the extraction function the current total sequence length in vLLM’s KV cache.

Metrics:

Relative L-inf: max|diff| / max|ref| — worst-case per-element error
Relative L2: ||diff||_2 / ||ref||_2 — overall magnitude of error
Bhattacharyya Coefficient (BC): Overlap between expected and actual error distributions (three-way only). BC≈1.0 means target errors match dtype-only errors; BC→0 means the target has a fundamentally different error pattern.

In two-way mode, ref is the expected KV. In three-way mode, ref is the baseline KV (ground truth), and both expected and actual errors are measured relative to it.

vLLM KV Cache Architecture#

Understanding vLLM’s KV cache design is essential for correct reconstruction.

Abstraction Hierarchy#

┌─────────────────────────────────────────────────────────────────────────┐
│                           MODEL LAYER LEVEL                            │
│  Each attention layer defines its own KVCacheSpec via                   │
│  get_kv_cache_spec()                                                   │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│                              KVCacheSpec                                │
│  Defines how ONE LAYER stores KV cache                                 │
│  Subclasses: FullAttentionSpec, SlidingWindowSpec, MLAAttentionSpec    │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│                           KVCacheGroupSpec                              │
│  Groups layers that SHARE THE SAME BLOCK TABLE                         │
│  - layer_names: list[str]  (e.g., ["layers.0.self_attn", ...])         │
│  - kv_cache_spec: KVCacheSpec (merged spec for the group)              │
└─────────────────────────────────────────────────────────────────────────┘
                                    │
                                    ▼
┌─────────────────────────────────────────────────────────────────────────┐
│                             KVCacheConfig                               │
│  Complete KV cache configuration for the model                         │
│  - kv_cache_groups: list[KVCacheGroupSpec]                             │
│  - num_blocks: int                                                     │
└─────────────────────────────────────────────────────────────────────────┘

Key Insight: Layers with different attention patterns (e.g., sliding window vs full attention) are placed in separate groups with separate block tables.

Model Examples#

Model	Layer Pattern	Groups	Block Tables
Llama 3.3 70B	80 full attention	1	1 shared
GPT-OSS 120B	18 sliding + 18 full (alternating)	2	2 separate
Gemma-3-27b	52 sliding + 10 full	7 (with padding)	7 separate

Paged Attention Format#

vLLM uses paged attention where KV cache is stored in fixed-size blocks:

# Paged format: [num_blocks, num_kv_heads, block_size, head_dim]
# Block table maps logical block index → physical block index

# To read token at position `pos`:
logical_block = pos // block_size
offset = pos % block_size
physical_block = block_table[logical_block]
k_value = k_cache[physical_block, :, offset, :]

Sliding Window Handling#

For models with sliding window attention (e.g., GPT-OSS alternating layers), there are two distinct vLLM code paths that affect KV cache layout:

Hybrid allocator disabled (current default for mixed full+sliding models):

vLLM’s unify_hybrid_kv_cache_specs() converts SlidingWindowSpec → FullAttentionSpec(sliding_window=N). All layers keep full block allocation. The spec_type in our config is FullAttentionSpec for both groups. Blocks are not freed — KV data is valid for all positions. The model still only attends to the window during computation, but the cache retains everything.

Hybrid allocator enabled (true SlidingWindowSpec):

Blocks outside the window are freed and replaced with null_block. KV data outside the window is garbage. spec_type is SlidingWindowSpec.

Note

The sliding_window field in the model config (e.g., 128) counts the current query token. HF’s past_key_values retains window - 1 positions (e.g., 127). The KVCacheSpec also stores the config value (128).

Reconstruction (reconstruct_contiguous_kv):

Only NaN-fills positions outside the window when spec_type == 'SlidingWindowSpec' (blocks actually freed). For FullAttentionSpec with sliding_window set, all positions are reconstructed normally.

Comparison (compare_kv_caches):

Requires all inputs to have the same sequence length per layer. Raises ValueError on shape mismatches — it does not silently align or pad.

For models with sliding window layers where HF golden KV is shorter than vLLM KV (HF trims past_key_values to the window), the caller must pad the shorter golden with NaN so both sides share the same per-layer sequence length before calling compare_kv_caches.

TP Sharding#

KV heads are sharded across tensor parallel ranks. Each rank’s cache has shape [num_blocks, heads_per_rank, block_size, head_dim] where heads_per_rank = max(1, total_num_kv_heads // tp_size) (vLLM convention).

Partitioned (total_num_kv_heads >= tp_size):
Each rank has unique heads. Concat all ranks on head dim.

Replicated (total_num_kv_heads < tp_size):
Multiple ranks share the same head. Stride by tp_size // total_num_kv_heads to pick one rank per replica group, then concat.

# Example: total=8, tp=8 → stride=1 → concat all 8 ranks → 8 heads
# Example: total=2, tp=8 → stride=4 → ranks [0,4] → concat → 2 heads
# Example: total=16, tp=8 → partitioned → concat all → 16 heads

The mode is determined by comparing kv_cache_config['total_num_kv_heads'] with kv_cache_config['tp_size'].

Reconstruction Algorithm#

Converting paged KV to contiguous format for comparison:

def reconstruct_contiguous_kv(paged_kv, kv_cache_config, block_tables, seq_len):
    # Build layer → group mapping from KVCacheConfig
    layer_to_group = {
        layer: group_idx
        for group_idx, group in enumerate(kv_cache_config['groups'])
        for layer in group['layer_names']
    }

    result = {}
    for layer_name, (k_paged, v_paged) in paged_kv.items():
        group_idx = layer_to_group[layer_name]
        block_size = kv_cache_config['groups'][group_idx]['block_size']
        block_table = block_tables[group_idx]

        # Gather physical blocks and reshape to contiguous
        num_blocks = ceil(seq_len / block_size)
        indices = block_table[:num_blocks]
        k_blocks = k_paged[indices]  # [num_blocks, heads, block_size, dim]
        k_cont = k_blocks.permute(1,0,2,3).reshape(heads, -1, dim)[:, :seq_len, :]

        # NaN-fill for SlidingWindowSpec (blocks freed)
        if spec_type == 'SlidingWindowSpec':
            k_cont[:, :valid_start, :] = float('nan')

        result[layer_name] = (k_cont, v_cont)

    return result

A batch variant reconstruct_contiguous_kv_batch reconstructs all sequences in a batch by calling reconstruct_contiguous_kv per sequence index.

Two-Way vs Three-Way Comparison#

Two-way (no baseline): compares expected vs actual directly.

result = compare_kv_caches(expected_kv, actual_kv)
# head.k_linf = |expected - actual| / |expected|

Three-way (with baseline): all errors are relative to baseline.

result, raw_errors = compare_kv_caches(
    expected_kv, actual_kv, baseline_kv=fp32_kv,
    return_raw_errors=True)

# head.k_linf     = |baseline - actual| / |baseline|    (target error)
# head.base_k_linf = |baseline - expected| / |baseline|  (dtype-only error)
# BC compares the two error distributions

If BC ≈ 1.0, the target’s errors look like dtype conversion errors (expected). If BC → 0, the target has a fundamentally different error pattern (bug).

Cross-Prompt Aggregation#

BC can be aggregated across multiple prompts to get a more robust signal. Errors are aligned by generation token index (0 = first decode token), so prompts with different lengths are properly aligned.

all_raw_errors = []  # from compare_kv_caches(..., return_raw_errors=True)
prompt_lens = [8, 12, 10]

agg = aggregate_kv_bc_across_prompts(all_raw_errors, prompt_lens)
# agg[gen_token_idx][layer_name].k_bc → aggregated BC

A compute_combined_bc_per_layer function computes a single BC per layer across all prompts and all tokens.

Integration with Logit Validation#

The recommended workflow captures KV caches during teacher-forced generation.

Note

extract_vllm_block_tables() automatically enables block table snapshotting on first call. Subsequent forward passes will preserve block tables even after requests are freed. There is zero overhead until the first call to get_block_tables().

from vllm_neuron.accuracy import logit_validation
from vllm_neuron.accuracy.kv_cache_analysis import (
    compare_kv_caches,
    extract_hf_kv_caches_teacher_forced,
    extract_vllm_kv_caches,
    extract_vllm_kv_cache_config,
    extract_vllm_block_tables,
    reconstruct_contiguous_kv,
)

kv_config = extract_vllm_kv_cache_config(llm)

def kv_extract_fn(seq_len):
    paged_kv = extract_vllm_kv_caches(llm, kv_config)
    block_tables = extract_vllm_block_tables(llm)
    return reconstruct_contiguous_kv(paged_kv, kv_config, block_tables, seq_len)

# logit_validation teacher-forces vLLM and captures KV at each step
passed, results, merged_kv = logit_validation(
    input_ids=input_ids.tolist(),
    generate_fn=vllm_generate_fn,
    expected_logits=expected_logits,
    baseline_logits=baseline_logits,  # None for two-way
    kv_extract_fn=kv_extract_fn,
)

# Compare with HF's KV caches
kv_result = compare_kv_caches(expected_kv, merged_kv, baseline_kv=baseline_kv)

Return types from logit_validation:

Two-way, no KV: passed (bool)
Two-way + KV: (passed, results, merged_kv)
Three-way, no KV: (passed, results)
Three-way + KV: (passed, results, merged_kv)

The results list contains per-token validation data including divergence information, which can be extracted and passed to the HTML visualization.

This ensures:

vLLM processes the exact same tokens as HF (teacher forcing)
KV caches are captured at the right points during generation
Multiple decode steps are merged into a single contiguous cache via _update_kv_state inside logit_validation

Incremental KV Merging#

logit_validation may call generate_fn multiple times (when divergence occurs, it re-feeds teacher tokens). _update_kv_state incrementally builds the merged KV cache:

First iteration: copies prompt KV + valid decode KV
Subsequent iterations: copies only newly validated decode KV
Divergence index determines how many decode tokens are valid per iteration

HF Golden Generation#

Two functions for generating HF KV caches:

generate_hf_logits_and_kv(model, input_ids, num_tokens): Autoregressive generation. Returns logits and KV. Performs an extra forward pass after the last generated token to ensure past_key_values includes KV for all positions up to and including the last generated token.
extract_hf_kv_caches_teacher_forced(model, input_ids, teacher_tokens): Single forward pass with the full sequence (prompt + teacher tokens). Optionally returns logits for the generated positions via return_logits=True.

Visualization#

Visualization helpers live in vllm_neuron.accuracy.kv_cache_visualize (the per-prompt report is export_html_report):

from vllm_neuron.accuracy.kv_cache_visualize import (
    export_html_report,
    export_aggregated_bc_html,
    launch_dashboard,
)

Per-prompt HTML heatmaps (Y=layer, X=token, color=max metric over heads):

export_html_report(result, "kv_prompt_0.html",
    prompt_len=prompt_len,
    divergence_indices=[8, 14])

Two-way mode shows:

L-inf and L2 error heatmaps using YlOrRd_r colorscale (yellow=low, red=high)
Default zmax=0.2 (20% relative error)

Three-way mode additionally shows:

Error ratio heatmaps (actual/expected, both vs baseline). Ratio ≈ 1.0x means actual error matches expected. Uses RdYlGn_r colorscale, zmax=3.0x.
BC heatmaps with custom colorscale focused on 0.8–1.0 range: 0.0–0.8 flat red, 0.8–0.85 orange, 0.85–0.90 yellow, 0.90–0.95 light green, 0.95–1.0 green. Below 0.8 is considered irrelevant.

Annotations on all heatmaps:

Red vertical line marks prefill/decode boundary
Orange dashed lines mark divergence points (where vLLM would have sampled a different token)

Aggregated BC HTML (across prompts):

agg = aggregate_kv_bc_across_prompts(all_raw_errors, prompt_lens)
export_aggregated_bc_html(agg, "aggregated_bc.html", num_prompts=3)

Uses the same BC-focused colorscale. Includes divergence count annotations.

Interactive dashboard (optional, requires dash):

from vllm_neuron.accuracy.kv_cache_visualize import launch_dashboard

launch_dashboard(result, port=8050)

API Reference#

Extraction#

# HF autoregressive (generates tokens, returns logits + KV)
logits, kv = generate_hf_logits_and_kv(model, input_ids, num_tokens)

# HF teacher-forced (single forward pass, optionally with logits)
kv = extract_hf_kv_caches_teacher_forced(model, input_ids, teacher_tokens)
logits, kv = extract_hf_kv_caches_teacher_forced(
    model, input_ids, teacher_tokens, return_logits=True)

# vLLM (via collective_rpc)
kv_config = extract_vllm_kv_cache_config(llm)
paged_kv = extract_vllm_kv_caches(llm, kv_config)
block_tables = extract_vllm_block_tables(llm)

Reconstruction#

# Single sequence
contiguous_kv = reconstruct_contiguous_kv(
    paged_kv, kv_config, block_tables, seq_len)
# Shape: [1, num_kv_heads, seq_len, head_dim]

# Full batch
batch_kv = reconstruct_contiguous_kv_batch(
    paged_kv, kv_config, block_tables, seq_len, batch_size)
# Shape: [batch_size, num_kv_heads, seq_len, head_dim]

Comparison#

# Two-way
result = compare_kv_caches(expected_kv, actual_kv)

# Three-way with raw errors for aggregation
result, raw_errors = compare_kv_caches(
    expected_kv, actual_kv, baseline_kv=fp32_kv,
    return_raw_errors=True)

# Per-token, per-layer, per-head metrics:
head = result[token_idx][layer_name][head_idx]
head.k_linf       # actual vs ref (ref = baseline if three-way, else expected)
head.base_k_linf  # expected vs baseline (three-way only)

# Per-token, per-layer BC (three-way only):
bc = result[token_idx][f"{layer_name}._bc"]
bc.k_bc  # BC between expected K errors and actual K errors

Data Structures#

@dataclass
class HeadMetrics:
    k_cos: float       # Cosine similarity (1.0 = identical)
    v_cos: float
    k_linf: float      # Relative L-inf: max|diff| / max|ref|
    v_linf: float
    k_l2: float        # Relative L2: ||diff||_2 / ||ref||_2
    v_l2: float
    base_k_linf: float  # expected vs baseline (three-way only)
    base_v_linf: float
    base_k_l2: float
    base_v_l2: float

@dataclass
class TokenKVMetrics:
    k_bc: float  # BC between expected and actual K errors
    v_bc: float  # BC between expected and actual V errors

@dataclass
class TokenKVErrors:
    base_k: np.ndarray  # [num_heads * head_dim] expected K errors vs baseline
    base_v: np.ndarray
    tgt_k: np.ndarray   # [num_heads * head_dim] actual K errors vs baseline
    tgt_v: np.ndarray

RPC Methods#

Added to NeuronModelRunner and NeuronWorker:

def get_kv_caches(self) -> dict[str, bytes]:
    """Returns {layer_name: torch.save({'k': k, 'v': v})}"""

def get_kv_cache_config(self) -> dict:
    """Returns KVCacheConfig as dict with groups, TP info, CP info."""

def get_block_tables(self) -> list[bytes]:
    """Returns [torch.save(block_table) for each group]."""

Dependencies#

Required: torch, numpy

Optional (visualization): plotly

Optional (interactive dashboard): dash, dash-bootstrap-components

Example Scripts#

examples/vllm_neuron/accuracy/run_kv_cache_analysis.py - Two-way and three-way KV cache comparison on Neuron

This document is relevant for: Trn2, Trn3

KV Cache Analysis Design

Contents

KV Cache Analysis Design#

Overview#

vLLM KV Cache Architecture#

Abstraction Hierarchy#

Model Examples#

Paged Attention Format#

Sliding Window Handling#

TP Sharding#

Reconstruction Algorithm#

Two-Way vs Three-Way Comparison#

Cross-Prompt Aggregation#

Integration with Logit Validation#

Incremental KV Merging#

HF Golden Generation#

Visualization#

API Reference#

Extraction#

Reconstruction#

Comparison#

Data Structures#

RPC Methods#

Dependencies#

Example Scripts#