This document is relevant for: Trn2, Trn3

Tutorial: Benchmark prefix caching with gpt-oss on vLLM on Neuron#

This tutorial guides you through benchmarking prefix caching with gpt-oss on vLLM on Neuron. When you have completed it, you will know how to measure whether prefix caching helps your workload and by how much, using your own data.

Overview#

Prefix caching reuses the KV cache computed for shared prompt prefixes across requests. When many requests share the same system prompt or context (RAG snippets, chat history), the server skips re-computing those tokens on subsequent requests, reducing time-to-first-token (TTFT).

Not every workload benefits equally. If your prompts rarely share a prefix, caching adds memory overhead without improving latency. This tutorial teaches you to answer “does prefix caching help my traffic?” by:

  1. Establishing a baseline without prefix caching.

  2. Enabling prefix caching and re-running the same workload.

  3. Comparing TTFT, throughput, and cache hit rate.

  4. Testing both a high-overlap workload (many requests share a prefix) and a no-overlap workload (unique prompts) to see when caching helps.

Before you start#

This tutorial assumes familiarity with:

Model details#

This tutorial uses gpt-oss 120B. The same methodology applies to gpt-oss 20B — adjust --tensor-parallel-size and the model identifier per the gpt-oss deployment tutorial.

Server settings for this tutorial:

  • tensor-parallel-size 8

  • max-num-seqs 4

  • max-model-len 8192

Limitations:

  • Prefix caching benefit scales with prefix length and reuse frequency. Workloads with short, unique prompts see no improvement.

  • The first request that establishes a prefix always pays full prefill cost.

Prerequisites#

  • Instance: A Trn2 or Trn3 instance.

  • Software: Neuron SDK 2.31 or later, vLLM on Neuron plugin (shipped with Neuron 2.31).

  • Model access: Access to openai/gpt-oss-120b on Hugging Face.

  • Datasets: The tutorial uses vLLM’s synthetic prefix_repetition dataset. It is generated by vllm bench serve directly, so there is nothing to download.


Prepare your environment#

Set environment variables required for gpt-oss compilation:

export VLLM_NEURON_COMPILATION_TIMEOUT=1200
export VLLM_EXECUTE_MODEL_TIMEOUT_SECONDS=1200
export NEURON_RT_ALLOW_LEGACY_NEFF=1  # Trn3 only

Step 1: Deploy the baseline server (no prefix caching)#

In this step, you will launch gpt-oss 120B with prefix caching explicitly disabled. This establishes the performance floor — every request computes its full prompt from scratch.

vllm serve openai/gpt-oss-120b \
    --tensor-parallel-size 8 \
    --max-num-seqs 4 \
    --max-model-len 8192 \
    --max-num-batched-tokens 1024 \
    --no-enable-prefix-caching \
    --hf-overrides '{"quantization_config": {}}' \
    --additional-config '{"neuron_config": {
        "num_batched_tokens_buckets": [1024],
        "num_seqs_buckets": [4],
        "kv_segment_size_buckets": [1024],
        "on_device_sampling_config": {"all_greedy": true}
    }}' \
    --port 8000

Wait for Application startup complete. in the server log before continuing. First-run compilation takes several minutes.

Confirm the server responds:

curl http://localhost:8000/health

Step 2: Benchmark the baseline with a high-overlap workload#

In this step, you will measure TTFT on a workload where requests share long prefixes. The prefix_repetition dataset builds each prompt as (one of N shared prefixes) + (unique suffix). With 50 prompts drawn from just 5 prefixes, each prefix is reused by 10 prompts on average — one cold request that computes the prefix, then nine warm requests that can reuse it.

vllm bench serve \
    --base-url http://localhost:8000 \
    --model openai/gpt-oss-120b \
    --dataset-name prefix_repetition \
    --prefix-repetition-prefix-len 3072 \
    --prefix-repetition-suffix-len 1024 \
    --prefix-repetition-num-prefixes 5 \
    --prefix-repetition-output-len 128 \
    --num-prompts 50 \
    --max-concurrency 1 \
    --ignore-eos \
    --save-result \
    --result-filename baseline_high_overlap.json

Key flags:

  • --prefix-repetition-prefix-len 3072 — the shared portion of each prompt is 3072 tokens, simulating a long shared system prompt. This is three times the 1024-token prefill segment (--max-num-batched-tokens), so the prefix spans three whole segments that prefix caching can skip on a warm request. If the prefix fit inside a single segment, no segment would be fully cached and there would be nothing to skip.

  • --prefix-repetition-suffix-len 1024 — the unique per-request tail (one segment). It is never cached, so it is always recomputed.

  • --prefix-repetition-num-prefixes 5 — 5 distinct prefixes shared across the 50 prompts. Fewer prefixes means heavier reuse.

  • --max-concurrency 1 — isolates prefill latency from queueing effects.

  • --ignore-eos — forces exactly --prefix-repetition-output-len output tokens for consistent measurement.

Record the median_ttft_ms and output_throughput values.

Step 3: Benchmark the baseline with a no-overlap workload#

In this step, you will measure the same metrics on a workload with zero prefix sharing. This shows what prefix caching cannot help.

Use the same dataset, but set --prefix-repetition-num-prefixes equal to --num-prompts. Every prompt then gets its own unique prefix, so no two requests share anything and caching has nothing to reuse:

vllm bench serve \
    --base-url http://localhost:8000 \
    --model openai/gpt-oss-120b \
    --dataset-name prefix_repetition \
    --prefix-repetition-prefix-len 3072 \
    --prefix-repetition-suffix-len 1024 \
    --prefix-repetition-num-prefixes 50 \
    --prefix-repetition-output-len 128 \
    --num-prompts 50 \
    --max-concurrency 1 \
    --ignore-eos \
    --save-result \
    --result-filename baseline_no_overlap.json

Changing only --prefix-repetition-num-prefixes (5 → 50) isolates prefix overlap as the single variable between the two workloads. Record median_ttft_ms and output_throughput.

Stop the baseline server with Ctrl+C.

Step 4: Deploy with prefix caching enabled#

In this step, you will relaunch gpt-oss with prefix caching enabled. The only change from the baseline is removing --no-enable-prefix-caching (prefix caching is on by default in vLLM). The kv_segment_size_buckets setting, which compiles the segmented-attention kernel that reads cached KV, is already present in the baseline command so both servers run the identical graph:

vllm serve openai/gpt-oss-120b \
    --tensor-parallel-size 8 \
    --max-num-seqs 4 \
    --max-model-len 8192 \
    --max-num-batched-tokens 1024 \
    --hf-overrides '{"quantization_config": {}}' \
    --additional-config '{"neuron_config": {
        "num_batched_tokens_buckets": [1024],
        "num_seqs_buckets": [4],
        "kv_segment_size_buckets": [1024],
        "on_device_sampling_config": {"all_greedy": true}
    }}' \
    --port 8000

Wait for Application startup complete. before continuing.

Step 5: Benchmark with prefix caching (high-overlap workload)#

Run the identical high-overlap benchmark. For each of the 5 prefixes, the first request computes the shared 3072-token prefix; the remaining requests that draw the same prefix reuse it from cache.

vllm bench serve \
    --base-url http://localhost:8000 \
    --model openai/gpt-oss-120b \
    --dataset-name prefix_repetition \
    --prefix-repetition-prefix-len 3072 \
    --prefix-repetition-suffix-len 1024 \
    --prefix-repetition-num-prefixes 5 \
    --prefix-repetition-output-len 128 \
    --num-prompts 50 \
    --max-concurrency 1 \
    --ignore-eos \
    --save-result \
    --result-filename cached_high_overlap.json

Step 6: Benchmark with prefix caching (no-overlap workload)#

Run the zero-overlap benchmark against the prefix-caching server:

vllm bench serve \
    --base-url http://localhost:8000 \
    --model openai/gpt-oss-120b \
    --dataset-name prefix_repetition \
    --prefix-repetition-prefix-len 3072 \
    --prefix-repetition-suffix-len 1024 \
    --prefix-repetition-num-prefixes 50 \
    --prefix-repetition-output-len 128 \
    --num-prompts 50 \
    --max-concurrency 1 \
    --ignore-eos \
    --save-result \
    --result-filename cached_no_overlap.json

Step 7: Compare results#

Compare all four runs side-by-side:

python3 - <<'PY'
import json

runs = {
    "baseline (high overlap)": "baseline_high_overlap.json",
    "cached (high overlap)":   "cached_high_overlap.json",
    "baseline (no overlap)":   "baseline_no_overlap.json",
    "cached (no overlap)":     "cached_no_overlap.json",
}

print(f"{'Run':<28} {'median_ttft_ms':>14} {'output_throughput':>18}")
print("-" * 64)
for label, path in runs.items():
    d = json.load(open(path))
    print(f"{label:<28} {d['median_ttft_ms']:>14.1f} {d['output_throughput']:>18.2f}")

# Compute speedups
bh = json.load(open("baseline_high_overlap.json"))
ch = json.load(open("cached_high_overlap.json"))
print(f"\nHigh-overlap TTFT reduction: "
      f"{(1 - ch['median_ttft_ms'] / bh['median_ttft_ms']) * 100:.1f}%")
print(f"High-overlap throughput gain: "
      f"{ch['output_throughput'] / bh['output_throughput']:.2f}x")
PY

Step 8: Run under sustained load#

Synthetic one-shot benchmarks at concurrency=1 hide queueing behavior. Real production traffic arrives concurrently, and prefix caching interacts with the scheduler’s batching decisions. This step measures performance under sustained concurrent load.

Run the high-overlap workload at concurrency matching --max-num-seqs:

# Against the prefix-caching server (still running from Step 4)
vllm bench serve \
    --base-url http://localhost:8000 \
    --model openai/gpt-oss-120b \
    --dataset-name prefix_repetition \
    --prefix-repetition-prefix-len 3072 \
    --prefix-repetition-suffix-len 1024 \
    --prefix-repetition-num-prefixes 5 \
    --prefix-repetition-output-len 128 \
    --num-prompts 200 \
    --max-concurrency 4 \
    --request-rate 2.0 \
    --ignore-eos \
    --save-result \
    --result-filename cached_sustained_c4.json

Key differences from the single-request benchmarks:

  • --max-concurrency 4 — matches max_num_seqs, keeping the decode batch full.

  • --request-rate 2.0 — sends 2 requests/second instead of back-to-back, simulating arrival patterns.

  • --num-prompts 200 — longer run for stable percentile estimates.

Compare the p50 and p99 TTFT against the concurrency-1 run to understand how batching and cache contention affect tail latency:

python3 - <<'PY'
import json

c1 = json.load(open("cached_high_overlap.json"))
c4 = json.load(open("cached_sustained_c4.json"))

print(f"{'Metric':<25} {'concurrency=1':>14} {'concurrency=4':>14}")
print("-" * 55)
for m in ("median_ttft_ms", "p99_ttft_ms", "output_throughput"):
    v1 = c1.get(m, 0)
    v4 = c4.get(m, 0)
    print(f"{m:<25} {v1:>14.1f} {v4:>14.1f}")
PY

Under sustained load, you should observe:

  • Higher total output_throughput (more requests in flight).

  • Slightly higher TTFT (prefill competes with decode for device time).

  • The TTFT benefit of prefix caching is preserved — cached requests still skip KV computation for the shared prefix.

Step 9: Observe cache hit rate#

The vLLM Prometheus endpoint exposes prefix caching statistics. While the server is running, query:

curl -s http://localhost:8000/metrics | grep -iE "prefix_cache|prompt_tokens_by_source"

Two complementary signals are available:

  • vllm:prefix_cache_queries / vllm:prefix_cache_hits — token-level counters. Their ratio is the fraction of prompt tokens served from cache.

  • vllm:prompt_tokens_by_source{source="local_cache_hit"} vs {source="local_compute"} — the same fraction, broken out by whether each prompt token was reused from cache or recomputed. This is the more direct signal: local_cache_hit / (local_cache_hit + local_compute).

Compute both:

python3 - <<'PY'
import re

import requests

metrics = requests.get("http://localhost:8000/metrics").text


def _val(pattern):
    m = re.search(pattern, metrics)
    return float(m.group(1)) if m else 0.0


queries = _val(r'vllm:prefix_cache_queries[^}]*}\s+([\d.]+)')
hits = _val(r'vllm:prefix_cache_hits[^}]*}\s+([\d.]+)')
if queries:
    print(f"prefix_cache hit rate: {hits / queries * 100:.1f}% "
          f"({int(hits)}/{int(queries)} tokens)")

cache_hit = _val(r'vllm:prompt_tokens_by_source\{[^}]*source="local_cache_hit"[^}]*\}\s+([\d.]+)')
compute = _val(r'vllm:prompt_tokens_by_source\{[^}]*source="local_compute"[^}]*\}\s+([\d.]+)')
total = cache_hit + compute
if total:
    print(f"prompt_tokens_by_source cache ratio: {cache_hit / total * 100:.1f}% "
          f"(cache_hit={int(cache_hit)}, compute={int(compute)})")
PY

Because prefix_repetition uses exact token counts, you can derive the expected hit rate before you run. With num_prefixes distinct prefixes and num_prompts total prompts, the first request per prefix is a cold miss (full compute) and the rest are warm hits on the prefix portion:

warm = num_prompts - num_prefixes          # requests that hit a warm prefix
cold = num_prefixes                          # one cold miss per prefix
cache_hit_tokens = warm * prefix_len
compute_tokens = cold * (prefix_len + suffix_len) + warm * suffix_len
expected_ratio = cache_hit_tokens / (cache_hit_tokens + compute_tokens)

For the high-overlap workload (num_prompts=50, num_prefixes=5, prefix_len=3072, suffix_len=1024) this is 45×3072 / (45×3072 + 5×4096 + 45×1024) 67.5%. Allow a small shortfall below this ceiling for block alignment and the cold-start phase.

For the no-overlap workload (num_prefixes=50), the expected ratio is 0% — every prefix is unique, so there is nothing to reuse.

Confirmation#

Prefix caching is working effectively when:

  • The high-overlap workload shows measurably lower median TTFT compared to the baseline. Warm requests skip the shared 3072-token prefix, ~67.5% of prompt tokens across the run.

  • The no-overlap workload shows no meaningful TTFT difference (confirms no regression from enabling caching).

  • The Prometheus cache hit rate matches the expected ratio derived from num_prefixes, num_prompts, prefix_len, and suffix_len.

  • Under sustained load, the throughput benefit is preserved.


Tuning kv_segment_size_buckets#

The main dial for prefix caching is kv_segment_size_buckets, which sets the compiled segment sizes for segmented attention — how the kernel reads cached KV.

Tune it when you change max_num_batched_tokens: the segment size must match the prefill bucket size (constraint: num_batched_tokens_buckets must equal kv_segment_size_buckets).

Supported values: [512], [1024], [2048], [4096].

Common issues#

  • No TTFT improvement on the high-overlap workload. The prefix_repetition dataset guarantees byte-identical prefixes, so this indicates a configuration or feature-enablement issue rather than a data issue. Check the Prometheus hit rate — if it’s 0%, confirm prefix caching is not disabled (--no-enable-prefix-caching must not be set). When using your own data, even a single whitespace difference between prompts prevents cache hits at the block boundary.

  • Cache evictions under sustained load. The cache is full and evicting entries before they can be reused. Reduce max_num_seqs or max_model_len to free memory for the prefix cache.

  • Server fails to start with OOM. The KV cache does not fit in HBM. Decrease VLLM_NEURON_KV_GMU_BUDGET_CAP_FRACTION (the fraction of the memory budget allocated to the KV cache) so the cache is sized smaller at startup.

Next steps#

This document is relevant for: Trn2, Trn3