| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Trace-replay with vLLM offline inference. |
| 4 | |
| 5 | Trace-replay lets you supply a known sequence of decode token IDs alongside |
| 6 | the prompt. Instead of sampling from the model distribution, the engine |
| 7 | injects each decode token deterministically, step by step. All other |
| 8 | outputs — logprobs, token ranks, text decoding — are computed faithfully |
| 9 | from the real logit distribution for that token. |
| 10 | |
| 11 | How it works: |
| 12 | 1. Start the engine with ``--enable-trace-replay`` (or |
| 13 | ``LLM(..., enable_trace_replay=True)``). |
| 14 | 2. Set ``SamplingParams.trace_decode_token_ids`` to the list of decode |
| 15 | token IDs you want to force. |
| 16 | 3. The engine will output exactly those tokens and stop. ``max_tokens`` |
| 17 | is overwritten with the trace length, and the trace is truncated if it |
| 18 | does not fit within ``max_model_len``. EOS tokens inside the trace |
| 19 | sequence do **not** halt generation early. |
| 20 | |
| 21 | Requires model runner V2. Requests are rejected with ``ValueError`` when |
| 22 | combined with any of: |
| 23 | * n > 1 * Speculative decoding |
| 24 | * prompt_logprobs * Structured outputs |
| 25 | * repetition_detection * thinking_token_budget |
| 26 | * bad_words |
| 27 | |
| 28 | Typical use-cases: |
| 29 | * Reproduce exact outputs from a previous run for benchmarking. |
| 30 | * Compute logprobs for an already-known output (e.g. reference answers). |
| 31 | * Dataset annotation: given (prompt, response) pairs, obtain per-token |
| 32 | logprob scores without altering the response. |
| 33 | |
| 34 | Usage: |
| 35 | python examples/generate/trace_replay_offline.py |
| 36 | python examples/generate/trace_replay_offline.py --model facebook/opt-125m |
| 37 | """ |
| 38 | |
| 39 | import argparse |
| 40 | |
| 41 | from vllm import LLM, SamplingParams |
| 42 | |
| 43 | DEFAULT_PROMPT = "Hello, my name is" |
| 44 | |
| 45 | |
| 46 | def build_llm(args: argparse.Namespace) -> LLM: |
| 47 | """Construct an LLM from common CLI args.""" |
| 48 | llm_kwargs: dict = { |
| 49 | "model": args.model, |
| 50 | "trust_remote_code": args.trust_remote_code, |
| 51 | "tensor_parallel_size": args.tensor_parallel_size, |
| 52 | "enforce_eager": args.enforce_eager, |
| 53 | "gpu_memory_utilization": args.gpu_memory_utilization, |
| 54 | "max_num_seqs": args.max_num_seqs, |
| 55 | "enable_trace_replay": True, |
| 56 | } |
| 57 | if args.max_model_len is not None: |
| 58 | llm_kwargs["max_model_len"] = args.max_model_len |
| 59 | return LLM(**llm_kwargs) |
| 60 | |
| 61 | |
| 62 | def run_normal_generation(llm: LLM, prompt: str, max_tokens: int) -> list[int]: |
| 63 | """Run a standard greedy generation and return the output token IDs.""" |
| 64 | sampling_params = SamplingParams( |
| 65 | temperature=0.0, |
| 66 | max_tokens=max_tokens, |
| 67 | # Request logprobs for the greedy token at each step so we can |
| 68 | # compare them against the trace-replay logprobs below. |
| 69 | logprobs=1, |
| 70 | ) |
| 71 | outputs = llm.generate([prompt], sampling_params=sampling_params) |
| 72 | result = outputs[0].outputs[0] |
| 73 | output_token_ids = list(result.token_ids) |
| 74 | print("[Normal generation]") |
| 75 | print(f" Prompt : {prompt!r}") |
| 76 | print(f" Output token IDs : {output_token_ids}") |
| 77 | print(f" Output text : {result.text!r}") |
| 78 | return output_token_ids |
| 79 | |
| 80 | |
| 81 | def run_trace_replay(llm: LLM, prompt: str, decode_token_ids: list[int]) -> None: |
| 82 | """Replay a known decode sequence and print per-token logprobs.""" |
| 83 | sampling_params = SamplingParams( |
| 84 | # Provide the decode tokens to replay. |
| 85 | trace_decode_token_ids=decode_token_ids, |
| 86 | # Request top-5 logprobs so we can inspect the distribution. |
| 87 | logprobs=5, |
| 88 | ) |
| 89 | outputs = llm.generate([prompt], sampling_params=sampling_params) |
| 90 | result = outputs[0].outputs[0] |
| 91 | replayed_ids = list(result.token_ids) |
| 92 | |
| 93 | print("\n[Trace-replay]") |
| 94 | print(f" Requested decode token IDs : {decode_token_ids}") |
| 95 | print(f" Replayed output token IDs : {replayed_ids}") |
| 96 | print(f" Replayed output text : {result.text!r}") |
| 97 | |
| 98 | # Verify the replayed tokens match exactly. |
| 99 | assert replayed_ids == decode_token_ids, ( |
| 100 | f"Mismatch!\n expected: {decode_token_ids}\n got: {replayed_ids}" |
| 101 | ) |
| 102 | print(" Replayed tokens match the requested trace exactly.") |
| 103 | |
| 104 | # Show per-token logprobs (computed from the real distribution). |
| 105 | if result.logprobs: |
| 106 | print("\n Per-token logprobs (trace token):") |
| 107 | for step, (token_id, logprob_dict) in enumerate( |
| 108 | zip(replayed_ids, result.logprobs) |
| 109 | ): |
| 110 | sampled_lp = logprob_dict.get(token_id) |
| 111 | lp_value = f"{sampled_lp.logprob:.4f}" if sampled_lp is not None else "n/a" |
| 112 | rank = sampled_lp.rank if sampled_lp is not None else "n/a" |
| 113 | print( |
| 114 | f" step {step:2d}: token_id={token_id:6d} " |
| 115 | f"logprob={lp_value} rank={rank}" |
| 116 | ) |
| 117 | |
| 118 | |
| 119 | def run_demo(args: argparse.Namespace) -> None: |
| 120 | print(f"Loading model: {args.model}\n") |
| 121 | llm = build_llm(args) |
| 122 | prompt = args.prompt |
| 123 | |
| 124 | print("=" * 60) |
| 125 | print("Step 1 — Normal greedy generation (captures decode tokens)") |
| 126 | print("=" * 60) |
| 127 | decode_token_ids = run_normal_generation(llm, prompt, max_tokens=8) |
| 128 | |
| 129 | print("\n" + "=" * 60) |
| 130 | print("Step 2 — Trace-replay with the captured tokens") |
| 131 | print("=" * 60) |
| 132 | run_trace_replay(llm, prompt, decode_token_ids) |
| 133 | |
| 134 | |
| 135 | def parse_args() -> argparse.Namespace: |
| 136 | parser = argparse.ArgumentParser( |
| 137 | description="Trace-replay with vLLM offline inference" |
| 138 | ) |
| 139 | parser.add_argument( |
| 140 | "--prompt", |
| 141 | type=str, |
| 142 | default=DEFAULT_PROMPT, |
| 143 | help="Text prompt to use for generation (default: %(default)r)", |
| 144 | ) |
| 145 | parser.add_argument( |
| 146 | "--model", |
| 147 | type=str, |
| 148 | default="facebook/opt-125m", |
| 149 | help="Name or path of the HuggingFace model to use", |
| 150 | ) |
| 151 | parser.add_argument( |
| 152 | "--trust-remote-code", |
| 153 | action="store_true", |
| 154 | help="Trust remote code from HuggingFace", |
| 155 | ) |
| 156 | parser.add_argument( |
| 157 | "--tensor-parallel-size", |
| 158 | type=int, |
| 159 | default=1, |
| 160 | help="Number of tensor parallel replicas", |
| 161 | ) |
| 162 | parser.add_argument( |
| 163 | "--enforce-eager", |
| 164 | action="store_true", |
| 165 | help="Always use eager-mode PyTorch (disable CUDA graph)", |
| 166 | ) |
| 167 | parser.add_argument( |
| 168 | "--gpu-memory-utilization", |
| 169 | type=float, |
| 170 | default=0.9, |
| 171 | help="Fraction of GPU memory to use", |
| 172 | ) |
| 173 | parser.add_argument( |
| 174 | "--max-model-len", |
| 175 | type=int, |
| 176 | default=None, |
| 177 | help="Model context length", |
| 178 | ) |
| 179 | parser.add_argument( |
| 180 | "--max-num-seqs", |
| 181 | type=int, |
| 182 | default=256, |
| 183 | help="Maximum number of sequences per iteration", |
| 184 | ) |
| 185 | return parser.parse_args() |
| 186 | |
| 187 | |
| 188 | def main() -> None: |
| 189 | args = parse_args() |
| 190 | run_demo(args) |
| 191 | |
| 192 | |
| 193 | if __name__ == "__main__": |
| 194 | main() |