@tduncan / Nvidia-Optimization-Stack-VLLM / commits / a1c5b1fd9f

[Core][V1] Support trace_decode_token_ids for deterministic decode replay (#46701)

Signed-off-by: ianzxliu <ianzxliu@tencent.com> Signed-off-by: aoshen02 <aoshen@inferact.ai> Signed-off-by: Zhixia Liu <liuzhixia.nju@gmail.com> Co-authored-by: ianzxliu <ianzxliu@tencent.com> Co-authored-by: Claude <noreply@anthropic.com> Co-authored-by: aoshen02 <aoshen@inferact.ai> Co-authored-by: Codex <noreply@openai.com> Co-authored-by: Nick Hill <nickhill123@gmail.com>

Zhixia Liu committed Aug 20, 2026 at 15:12 UTC a1c5b1fd9f6ef06a4fa236b7d48350115e5688b9
13 files changed +931 -9
docs/serving/online_serving/trace_replay.md new
+60
@@ -0,0 +1,60 @@
1 +# Trace Replay
2 +
3 +Trace replay forces the engine to emit a predetermined token sequence during decoding while computing real logprobs from the model's unmodified logit distribution. The primary use case is comparing logprob distributions across different configurations:
4 +
5 +- **Inference config diff**: compare logprobs between different quantization schemes, tensor parallelism layouts, or attention backends for the same model.
6 +- **Train vs inference diff**: replay a training-time token sequence through the inference engine to detect logprob divergence caused by numerical differences between training and serving frameworks.
7 +
8 +## Requirements
9 +
10 +Trace replay reserves a per-request trace buffer, so it is off by default. Enable it with `--enable-trace-replay` (or `LLM(..., enable_trace_replay=True)`). It is only supported by model runner V2; requests are rejected when either condition is not met.
11 +
12 +## Usage
13 +
14 +```python
15 +from vllm import LLM, SamplingParams
16 +
17 +llm = LLM(model="Qwen/Qwen3-0.6B", enable_trace_replay=True)
18 +
19 +# Token sequence captured from a previous run or training log
20 +trace_tokens = [15, 284, 1026, 374]
21 +
22 +params = SamplingParams(
23 + trace_decode_token_ids=trace_tokens,
24 + logprobs=5,
25 +)
26 +outputs = llm.generate(["Once upon a time"], sampling_params=params)
27 +
28 +for token, logprob in zip(
29 + outputs[0].outputs[0].token_ids,
30 + outputs[0].outputs[0].logprobs,
31 +):
32 + print(f"token={token} logprob={logprob[token].logprob:.4f}")
33 +```
34 +
35 +The output tokens will always be `[15, 284, 1026, 374]`. The logprobs reflect the model's true probability for each forced token under the current inference configuration.
36 +
37 +To compare configurations, run the same trace against two engine setups and diff the per-token logprobs.
38 +
39 +## Behavior
40 +
41 +When `trace_decode_token_ids` is set:
42 +
43 +- `max_tokens` is automatically set to the trace length.
44 +- The trace is truncated if it does not fit within `max_model_len` given the prompt length.
45 +- All stop conditions (EOS, stop strings, stop token IDs) are disabled.
46 +- Generation produces exactly the trace tokens, then stops.
47 +
48 +## Limitations
49 +
50 +`trace_decode_token_ids` is incompatible with the following features (raises `ValueError`):
51 +
52 +| Feature | Reason |
53 +| --- | --- |
54 +| `n > 1` | Trace replay produces a single deterministic sequence |
55 +| `prompt_logprobs` | Expanded logit layout conflicts with trace kernel indexing |
56 +| Speculative decoding | Multi-token speculation conflicts with single-token trace stepping |
57 +| Structured outputs | Grammar constraints conflict with forced token injection |
58 +| `repetition_detection` | Would terminate the request on repeated patterns in the trace |
59 +| `thinking_token_budget` | Logit masking corrupts logprobs for trace tokens |
60 +| `bad_words` | Logit masking corrupts logprobs for trace tokens |
examples/generate/trace_replay_offline.py new
+195
@@ -0,0 +1,195 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""
4 +Trace-replay with vLLM offline inference.
5 +
6 +Trace-replay lets you supply a known sequence of decode token IDs alongside
7 +the prompt. Instead of sampling from the model distribution, the engine
8 +injects each decode token deterministically, step by step. All other
9 +outputs — logprobs, token ranks, text decoding — are computed faithfully
10 +from the real logit distribution for that token.
11 +
12 +How it works:
13 + 1. Start the engine with ``--enable-trace-replay`` (or
14 + ``LLM(..., enable_trace_replay=True)``).
15 + 2. Set ``SamplingParams.trace_decode_token_ids`` to the list of decode
16 + token IDs you want to force.
17 + 3. The engine will output exactly those tokens and stop. ``max_tokens``
18 + is overwritten with the trace length, and the trace is truncated if it
19 + does not fit within ``max_model_len``. EOS tokens inside the trace
20 + sequence do **not** halt generation early.
21 +
22 +Requires model runner V2. Requests are rejected with ``ValueError`` when
23 +combined with any of:
24 + * n > 1 * Speculative decoding
25 + * prompt_logprobs * Structured outputs
26 + * repetition_detection * thinking_token_budget
27 + * bad_words
28 +
29 +Typical use-cases:
30 + * Reproduce exact outputs from a previous run for benchmarking.
31 + * Compute logprobs for an already-known output (e.g. reference answers).
32 + * Dataset annotation: given (prompt, response) pairs, obtain per-token
33 + logprob scores without altering the response.
34 +
35 +Usage:
36 + python examples/generate/trace_replay_offline.py
37 + python examples/generate/trace_replay_offline.py --model facebook/opt-125m
38 +"""
39 +
40 +import argparse
41 +
42 +from vllm import LLM, SamplingParams
43 +
44 +DEFAULT_PROMPT = "Hello, my name is"
45 +
46 +
47 +def build_llm(args: argparse.Namespace) -> LLM:
48 + """Construct an LLM from common CLI args."""
49 + llm_kwargs: dict = {
50 + "model": args.model,
51 + "trust_remote_code": args.trust_remote_code,
52 + "tensor_parallel_size": args.tensor_parallel_size,
53 + "enforce_eager": args.enforce_eager,
54 + "gpu_memory_utilization": args.gpu_memory_utilization,
55 + "max_num_seqs": args.max_num_seqs,
56 + "enable_trace_replay": True,
57 + }
58 + if args.max_model_len is not None:
59 + llm_kwargs["max_model_len"] = args.max_model_len
60 + return LLM(**llm_kwargs)
61 +
62 +
63 +def run_normal_generation(llm: LLM, prompt: str, max_tokens: int) -> list[int]:
64 + """Run a standard greedy generation and return the output token IDs."""
65 + sampling_params = SamplingParams(
66 + temperature=0.0,
67 + max_tokens=max_tokens,
68 + # Request logprobs for the greedy token at each step so we can
69 + # compare them against the trace-replay logprobs below.
70 + logprobs=1,
71 + )
72 + outputs = llm.generate([prompt], sampling_params=sampling_params)
73 + result = outputs[0].outputs[0]
74 + output_token_ids = list(result.token_ids)
75 + print("[Normal generation]")
76 + print(f" Prompt : {prompt!r}")
77 + print(f" Output token IDs : {output_token_ids}")
78 + print(f" Output text : {result.text!r}")
79 + return output_token_ids
80 +
81 +
82 +def run_trace_replay(llm: LLM, prompt: str, decode_token_ids: list[int]) -> None:
83 + """Replay a known decode sequence and print per-token logprobs."""
84 + sampling_params = SamplingParams(
85 + # Provide the decode tokens to replay.
86 + trace_decode_token_ids=decode_token_ids,
87 + # Request top-5 logprobs so we can inspect the distribution.
88 + logprobs=5,
89 + )
90 + outputs = llm.generate([prompt], sampling_params=sampling_params)
91 + result = outputs[0].outputs[0]
92 + replayed_ids = list(result.token_ids)
93 +
94 + print("\n[Trace-replay]")
95 + print(f" Requested decode token IDs : {decode_token_ids}")
96 + print(f" Replayed output token IDs : {replayed_ids}")
97 + print(f" Replayed output text : {result.text!r}")
98 +
99 + # Verify the replayed tokens match exactly.
100 + assert replayed_ids == decode_token_ids, (
101 + f"Mismatch!\n expected: {decode_token_ids}\n got: {replayed_ids}"
102 + )
103 + print(" Replayed tokens match the requested trace exactly.")
104 +
105 + # Show per-token logprobs (computed from the real distribution).
106 + if result.logprobs:
107 + print("\n Per-token logprobs (trace token):")
108 + for step, (token_id, logprob_dict) in enumerate(
109 + zip(replayed_ids, result.logprobs)
110 + ):
111 + sampled_lp = logprob_dict.get(token_id)
112 + lp_value = f"{sampled_lp.logprob:.4f}" if sampled_lp is not None else "n/a"
113 + rank = sampled_lp.rank if sampled_lp is not None else "n/a"
114 + print(
115 + f" step {step:2d}: token_id={token_id:6d} "
116 + f"logprob={lp_value} rank={rank}"
117 + )
118 +
119 +
120 +def run_demo(args: argparse.Namespace) -> None:
121 + print(f"Loading model: {args.model}\n")
122 + llm = build_llm(args)
123 + prompt = args.prompt
124 +
125 + print("=" * 60)
126 + print("Step 1 — Normal greedy generation (captures decode tokens)")
127 + print("=" * 60)
128 + decode_token_ids = run_normal_generation(llm, prompt, max_tokens=8)
129 +
130 + print("\n" + "=" * 60)
131 + print("Step 2 — Trace-replay with the captured tokens")
132 + print("=" * 60)
133 + run_trace_replay(llm, prompt, decode_token_ids)
134 +
135 +
136 +def parse_args() -> argparse.Namespace:
137 + parser = argparse.ArgumentParser(
138 + description="Trace-replay with vLLM offline inference"
139 + )
140 + parser.add_argument(
141 + "--prompt",
142 + type=str,
143 + default=DEFAULT_PROMPT,
144 + help="Text prompt to use for generation (default: %(default)r)",
145 + )
146 + parser.add_argument(
147 + "--model",
148 + type=str,
149 + default="facebook/opt-125m",
150 + help="Name or path of the HuggingFace model to use",
151 + )
152 + parser.add_argument(
153 + "--trust-remote-code",
154 + action="store_true",
155 + help="Trust remote code from HuggingFace",
156 + )
157 + parser.add_argument(
158 + "--tensor-parallel-size",
159 + type=int,
160 + default=1,
161 + help="Number of tensor parallel replicas",
162 + )
163 + parser.add_argument(
164 + "--enforce-eager",
165 + action="store_true",
166 + help="Always use eager-mode PyTorch (disable CUDA graph)",
167 + )
168 + parser.add_argument(
169 + "--gpu-memory-utilization",
170 + type=float,
171 + default=0.9,
172 + help="Fraction of GPU memory to use",
173 + )
174 + parser.add_argument(
175 + "--max-model-len",
176 + type=int,
177 + default=None,
178 + help="Model context length",
179 + )
180 + parser.add_argument(
181 + "--max-num-seqs",
182 + type=int,
183 + default=256,
184 + help="Maximum number of sequences per iteration",
185 + )
186 + return parser.parse_args()
187 +
188 +
189 +def main() -> None:
190 + args = parse_args()
191 + run_demo(args)
192 +
193 +
194 +if __name__ == "__main__":
195 + main()
tests/v1/engine/test_input_processor_trace_replay.py new
+110
@@ -0,0 +1,110 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import time
5 +from types import SimpleNamespace
6 +from unittest.mock import patch
7 +
8 +import pytest
9 +
10 +from vllm import SamplingParams
11 +from vllm.config import VllmConfig
12 +from vllm.exceptions import VLLMValidationError
13 +from vllm.v1.core.sched.utils import check_stop
14 +from vllm.v1.engine.input_processor import InputProcessor
15 +from vllm.v1.request import Request, RequestStatus
16 +
17 +pytestmark = pytest.mark.skip_global_cleanup
18 +
19 +
20 +def _make_request(params: SamplingParams) -> Request:
21 + return Request(
22 + request_id="req-0",
23 + client_index=0,
24 + prompt_token_ids=[1, 2, 3],
25 + sampling_params=params,
26 + pooling_params=None,
27 + arrival_time=time.time(),
28 + )
29 +
30 +
31 +def _normalize(params: SamplingParams, prompt_len: int, max_model_len: int) -> None:
32 + processor = SimpleNamespace(
33 + model_config=SimpleNamespace(max_model_len=max_model_len)
34 + )
35 + InputProcessor._normalize_trace_replay_params(processor, params, prompt_len)
36 +
37 +
38 +def test_normalize_trace_replay_params():
39 + params = SamplingParams(
40 + max_tokens=16,
41 + min_tokens=10,
42 + stop_token_ids=[20],
43 + trace_decode_token_ids=[10, 20, 30],
44 + )
45 + params.update_from_generation_config({}, eos_token_id=20)
46 +
47 + _normalize(params, prompt_len=3, max_model_len=128)
48 +
49 + assert params.max_tokens == 3
50 + assert params.min_tokens == 0
51 + assert params.ignore_eos is True
52 + assert params.eos_token_id is None
53 + assert params.stop_token_ids == []
54 + assert params.all_stop_token_ids == set()
55 +
56 + request = _make_request(params)
57 + assert request.max_tokens == 3
58 + request.append_output_token_ids([10, 20, 30])
59 + assert check_stop(request, max_model_len=128)
60 + assert request.status == RequestStatus.FINISHED_LENGTH_CAPPED
61 +
62 +
63 +def test_trace_longer_than_remaining_context_is_truncated():
64 + """The trace is staged into a max_model_len-wide row, so it must be cut.
65 +
66 + An explicitly set max_tokens is never clamped to max_model_len, so it cannot
67 + bound the staged write.
68 + """
69 + params = SamplingParams(max_tokens=100, trace_decode_token_ids=list(range(20)))
70 +
71 + _normalize(params, prompt_len=6, max_model_len=10)
72 +
73 + assert params.trace_decode_token_ids == [0, 1, 2, 3]
74 + assert params.max_tokens == 4
75 +
76 +
77 +def _validate(enable_trace_replay: bool) -> None:
78 + """Run _validate_params' trace gating with the rest of verify() stubbed."""
79 + processor = SimpleNamespace(
80 + model_config=SimpleNamespace(
81 + return_sampling_mask=False,
82 + enable_trace_replay=enable_trace_replay,
83 + ),
84 + vllm_config=SimpleNamespace(reasoning_config=None),
85 + speculative_config=None,
86 + structured_outputs_config=None,
87 + tokenizer=None,
88 + )
89 + params = SamplingParams(trace_decode_token_ids=[1, 2, 3])
90 + with patch.object(SamplingParams, "verify"):
91 + InputProcessor._validate_params(processor, params, ("generate",))
92 +
93 +
94 +def test_trace_request_rejected_when_feature_disabled():
95 + with pytest.raises(VLLMValidationError, match="--enable-trace-replay"):
96 + _validate(enable_trace_replay=False)
97 +
98 +
99 +def test_trace_request_accepted_when_feature_enabled():
100 + _validate(enable_trace_replay=True)
101 +
102 +
103 +def test_trace_replay_requires_v2_model_runner():
104 + config = SimpleNamespace(
105 + model_config=SimpleNamespace(enable_trace_replay=True),
106 + use_v2_model_runner=False,
107 + )
108 +
109 + with pytest.raises(ValueError, match="trace replay requires Model Runner V2"):
110 + VllmConfig._verify_trace_replay_config(config)
tests/v1/sample/test_trace_replay_params.py new
+111
@@ -0,0 +1,111 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +import pytest
5 +
6 +from vllm import SamplingParams
7 +from vllm.exceptions import VLLMValidationError
8 +from vllm.sampling_params import StructuredOutputsParams
9 +
10 +pytestmark = pytest.mark.skip_global_cleanup
11 +
12 +# ---------------------------------------------------------------------------
13 +# SamplingParams field
14 +# ---------------------------------------------------------------------------
15 +
16 +
17 +def test_sampling_params_trace_field_defaults_to_none():
18 + params = SamplingParams(max_tokens=10)
19 + assert params.trace_decode_token_ids is None
20 +
21 +
22 +def test_sampling_params_trace_field_accepts_list():
23 + ids = [100, 200, 300]
24 + params = SamplingParams(trace_decode_token_ids=ids)
25 + assert params.trace_decode_token_ids == ids
26 +
27 +
28 +def test_sampling_params_trace_field_preserved_by_clone():
29 + ids = [1, 2, 3]
30 + params = SamplingParams(trace_decode_token_ids=ids)
31 + cloned = params.clone()
32 + assert cloned.trace_decode_token_ids == ids
33 + assert cloned.trace_decode_token_ids is not params.trace_decode_token_ids
34 +
35 +
36 +def test_sampling_params_trace_field_rejects_empty_list():
37 + params = SamplingParams(trace_decode_token_ids=[])
38 + with pytest.raises(ValueError, match="non-empty"):
39 + params._validate_trace_replay(
40 + _make_model_config(vocab_size=100), speculative_config=None
41 + )
42 +
43 +
44 +def test_sampling_params_trace_field_requires_single_output():
45 + params = SamplingParams(n=2, trace_decode_token_ids=[1])
46 + with pytest.raises(ValueError, match="requires n=1"):
47 + params._validate_trace_replay(
48 + _make_model_config(vocab_size=100), speculative_config=None
49 + )
50 +
51 +
52 +@pytest.mark.parametrize("invalid_ids", [[-1, 5], [1, "2"]])
53 +def test_sampling_params_trace_field_rejects_invalid_token_ids(invalid_ids):
54 + params = SamplingParams(trace_decode_token_ids=invalid_ids)
55 + with pytest.raises(ValueError, match="non-negative integers"):
56 + params._validate_trace_replay(
57 + _make_model_config(vocab_size=100), speculative_config=None
58 + )
59 +
60 +
61 +def _make_model_config(vocab_size: int):
62 + from unittest.mock import Mock
63 +
64 + model_config = Mock()
65 + model_config.get_vocab_size = lambda: vocab_size
66 + return model_config
67 +
68 +
69 +def test_validate_trace_replay_accepts_in_vocab():
70 + params = SamplingParams(trace_decode_token_ids=[0, 50, 99])
71 + # Should not raise.
72 + params._validate_trace_replay(
73 + _make_model_config(vocab_size=100), speculative_config=None
74 + )
75 +
76 +
77 +def test_validate_trace_replay_rejects_out_of_vocab():
78 + # The non-negative check passes at construction, but the token id exceeds
79 + # the vocabulary; verify() must reject it before it reaches the sampler.
80 + params = SamplingParams(trace_decode_token_ids=[0, 100])
81 + with pytest.raises(VLLMValidationError, match="out-of-vocab"):
82 + params._validate_trace_replay(
83 + _make_model_config(vocab_size=100), speculative_config=None
84 + )
85 +
86 +
87 +def test_validate_trace_replay_noop_when_unset():
88 + params = SamplingParams(max_tokens=4)
89 + # Should not raise when the field is unset.
90 + params._validate_trace_replay(
91 + _make_model_config(vocab_size=100), speculative_config=None
92 + )
93 +
94 +
95 +def test_trace_decode_token_ids_rejects_speculative_decoding():
96 + params = SamplingParams(trace_decode_token_ids=[1])
97 + with pytest.raises(ValueError, match="not supported with speculative decoding"):
98 + params._validate_trace_replay(
99 + _make_model_config(vocab_size=100), speculative_config=object()
100 + )
101 +
102 +
103 +def test_trace_decode_token_ids_rejects_structured_outputs():
104 + params = SamplingParams(
105 + trace_decode_token_ids=[1],
106 + structured_outputs=StructuredOutputsParams(json_object=True),
107 + )
108 + with pytest.raises(ValueError, match="not supported with structured outputs"):
109 + params._validate_trace_replay(
110 + _make_model_config(vocab_size=100), speculative_config=None
111 + )
tests/v1/worker/test_gpu_trace_replay.py new
+195
@@ -0,0 +1,195 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""Tests for Model Runner V2 inference trace-replay.
4 +
5 +Trace-replay forces the sampler to emit a predetermined sequence of decode
6 +token IDs. The replay step for a request is derived from GPU state as
7 +``total_len - prompt_len`` (output tokens produced so far), and the trace token
8 +overwrites the sampled token in place before logprobs are computed.
9 +"""
10 +
11 +from types import SimpleNamespace
12 +from typing import cast
13 +
14 +import pytest
15 +import torch
16 +
17 +pytest.importorskip("triton")
18 +if not torch.cuda.is_available():
19 + pytest.skip("CUDA required for trace-replay tests", allow_module_level=True)
20 +
21 +from vllm.sampling_params import SamplingParams
22 +from vllm.v1.worker.gpu.sample.trace_replay import (
23 + TraceReplayState,
24 + apply_trace_tokens,
25 +)
26 +from vllm.v1.worker.gpu.states import RequestState
27 +
28 +DEVICE = "cuda"
29 +TEST_MAX_MODEL_LEN = 2048
30 +
31 +
32 +def _i32(x) -> torch.Tensor:
33 + return torch.tensor(x, dtype=torch.int32, device=DEVICE)
34 +
35 +
36 +def _i64(x) -> torch.Tensor:
37 + return torch.tensor(x, dtype=torch.int64, device=DEVICE)
38 +
39 +
40 +def _trace_state(max_num_reqs: int) -> TraceReplayState:
41 + """Build a state whose req_states exposes the buffers apply_trace reads.
42 +
43 + total_len/prompt_len are written by the caller to position each request at
44 + the desired replay step.
45 + """
46 + req_states = cast(
47 + RequestState,
48 + SimpleNamespace(
49 + max_num_reqs=max_num_reqs,
50 + max_model_len=TEST_MAX_MODEL_LEN,
51 + device=torch.device(DEVICE),
52 + total_len=SimpleNamespace(
53 + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE)
54 + ),
55 + prompt_len=SimpleNamespace(
56 + gpu=torch.zeros(max_num_reqs, dtype=torch.int32, device=DEVICE)
57 + ),
58 + ),
59 + )
60 + return TraceReplayState(req_states)
61 +
62 +
63 +def _set_lens(state: TraceReplayState, total_len, prompt_len) -> None:
64 + """Position each request at a replay step via total_len - prompt_len."""
65 + state.req_states.total_len.gpu[: len(total_len)] = _i32(total_len)
66 + state.req_states.prompt_len.gpu[: len(prompt_len)] = _i32(prompt_len)
67 +
68 +
69 +# ------------------------------ Kernel ------------------------------------
70 +
71 +
72 +def test_replay_overwrites_sampled_at_each_step():
73 + """The trace token for the current step replaces the sampled token."""
74 + trace = [[100, 101, 102], [200, 201, 202]]
75 + trace_len = _i32([3, 3])
76 + trace_token_ids = torch.zeros(
77 + 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
78 + )
79 + for i, t in enumerate(trace):
80 + trace_token_ids[i, : len(t)] = _i32(t)
81 + prompt_len = _i32([5, 8])
82 + idx_mapping = _i32([0, 1])
83 +
84 + for step in range(3):
85 + sampled = _i64([-7, -7]) # sentinel that must be overwritten
86 + total_len = _i32([5 + step, 8 + step])
87 + apply_trace_tokens(
88 + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len
89 + )
90 + assert sampled.tolist() == [trace[0][step], trace[1][step]]
91 +
92 +
93 +def test_past_end_of_trace_leaves_sampled_untouched():
94 + """Once step >= trace_len, the sampler's own token is kept."""
95 + trace_token_ids = torch.zeros(
96 + 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
97 + )
98 + trace_token_ids[0, :2] = _i32([100, 101])
99 + trace_len = _i32([2])
100 + prompt_len = _i32([4])
101 + idx_mapping = _i32([0])
102 +
103 + sampled = _i64([999])
104 + total_len = _i32([4 + 2]) # step == 2 == trace_len -> out of range
105 + apply_trace_tokens(
106 + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len
107 + )
108 + assert sampled.tolist() == [999]
109 +
110 +
111 +def test_non_trace_request_untouched():
112 + """trace_len == 0 means the request never uses replay."""
113 + trace_token_ids = torch.zeros(
114 + 1, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
115 + )
116 + trace_len = _i32([0])
117 + prompt_len = _i32([3])
118 + idx_mapping = _i32([0])
119 +
120 + sampled = _i64([42])
121 + total_len = _i32([3])
122 + apply_trace_tokens(
123 + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len
124 + )
125 + assert sampled.tolist() == [42]
126 +
127 +
128 +def test_idx_mapping_indirection_and_negative_skip():
129 + """batch_idx -> req_state_idx indirection, and negative entries are skipped."""
130 + # req_state 0: no trace; req_state 1: trace [500, 501].
131 + trace_token_ids = torch.zeros(
132 + 2, TEST_MAX_MODEL_LEN, dtype=torch.int32, device=DEVICE
133 + )
134 + trace_token_ids[1, :2] = _i32([500, 501])
135 + trace_len = _i32([0, 2])
136 + prompt_len = _i32([10, 6])
137 +
138 + # batch: [maps to state 1, masked (-1), maps to state 0].
139 + idx_mapping = _i32([1, -1, 0])
140 + sampled = _i64([1, 2, 3])
141 + # Indexed by req_state_idx: state 0 at total_len=10, state 1 at step 1.
142 + total_len = _i32([10, 6 + 1])
143 + apply_trace_tokens(
144 + sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len
145 + )
146 + # batch 0 -> state 1 step 1 -> 501; batch 1 masked; batch 2 -> state 0 no trace.
147 + assert sampled.tolist() == [501, 2, 3]
148 +
149 +
150 +# --------------------------- TraceReplayState -----------------------------
151 +
152 +
153 +def test_state_end_to_end():
154 + """add_request -> apply_staged_writes -> apply_trace overwrites correctly."""
155 + state = _trace_state(4)
156 + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22, 33]))
157 + state.add_request(1, SamplingParams()) # no trace
158 + state.apply_staged_writes()
159 +
160 + idx_mapping = _i32([0, 1])
161 + sampled = _i64([-1, -1])
162 + _set_lens(state, [7 + 1, 7, 0, 0], [7, 7, 0, 0]) # req 0 at step 1
163 + state.apply_trace(sampled, idx_mapping)
164 + assert sampled.tolist() == [22, -1]
165 +
166 +
167 +def test_state_leaves_non_trace_batch_unchanged():
168 + """Requests with trace_len == 0 remain unchanged after a trace was seen."""
169 + state = _trace_state(4)
170 + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22]))
171 + state.add_request(1, SamplingParams())
172 + state.apply_staged_writes()
173 +
174 + # Only the non-trace request (state 1) is in this batch.
175 + idx_mapping = _i32([1])
176 + sampled = _i64([555])
177 + _set_lens(state, [0, 7, 0, 0], [7, 7, 0, 0])
178 + state.apply_trace(sampled, idx_mapping)
179 + assert sampled.tolist() == [555]
180 +
181 +
182 +def test_slot_reuse_clears_trace():
183 + """Reusing a slot for a non-trace request must not replay stale tokens."""
184 + state = _trace_state(2)
185 + state.add_request(0, SamplingParams(trace_decode_token_ids=[11, 22]))
186 + state.apply_staged_writes()
187 + # Slot 0 reused by a request without a trace.
188 + state.add_request(0, SamplingParams())
189 + state.apply_staged_writes()
190 +
191 + idx_mapping = _i32([0])
192 + sampled = _i64([888])
193 + _set_lens(state, [3, 0], [3, 0])
194 + state.apply_trace(sampled, idx_mapping)
195 + assert sampled.tolist() == [888]
vllm/config/model.py
+7
@@ -262,6 +262,12 @@ class ModelConfig:
262 equivalent exponential-race sampling. FP64 preserves lower-tail sampling
263 events that fp32 uniform/exponential draws can truncate, at the cost of
264 significantly lower throughput on most GPUs."""
265 + enable_trace_replay: bool = False
266 + """Whether to allow requests to set
267 + `SamplingParams.trace_decode_token_ids`, which forces decoding to follow a
268 + predetermined token sequence while still computing real logprobs. Reserved
269 + for debugging and RL workflows: enabling it reserves a per-request trace
270 + buffer, so it is off by default."""
271 disable_sliding_window: bool = False
272 """Whether to disable sliding window. If True, we will disable the sliding
273 window functionality of the model, capping to sliding window size. If the
@@ -424,6 +430,7 @@ class ModelConfig:
430 "return_sampling_mask",
431 "logprobs_mode",
432 "use_fp64_gumbel",
433 + "enable_trace_replay",
434 "disable_cascade_attn",
435 "skip_tokenizer_init",
436 "served_model_name",
vllm/config/vllm.py
+8
@@ -1091,6 +1091,13 @@ class VllmConfig:
1091 "are normalized over the same nucleus as the sampling mask"
1092 )
1093
1094 + def _verify_trace_replay_config(self) -> None:
1095 + model_config = self.model_config
1096 + if model_config is None or not model_config.enable_trace_replay:
1097 + return
1098 + if not self.use_v2_model_runner:
1099 + raise ValueError("trace replay requires Model Runner V2")
1100 +
1101 def __post_init__(self):
1102 """Verify configs are valid & consistent with each other."""
1103
@@ -1141,6 +1148,7 @@ class VllmConfig:
1148 )
1149
1150 self._verify_sampling_replay_config()
1151 + self._verify_trace_replay_config()
1152
1153 if self.lora_config is not None:
1154 self.lora_config.verify_with_model_config(self.model_config)
vllm/engine/arg_utils.py
+5
@@ -543,6 +543,7 @@ class EngineArgs:
543 max_logprobs: int = ModelConfig.max_logprobs
544 logprobs_mode: LogprobsMode = ModelConfig.logprobs_mode
545 use_fp64_gumbel: bool = ModelConfig.use_fp64_gumbel
546 + enable_trace_replay: bool = ModelConfig.enable_trace_replay
547 disable_log_stats: bool = False
548 aggregate_engine_logging: bool = False
549 revision: str | None = ModelConfig.revision
@@ -884,6 +885,9 @@ class EngineArgs:
885 model_group.add_argument("--max-logprobs", **model_kwargs["max_logprobs"])
886 model_group.add_argument("--logprobs-mode", **model_kwargs["logprobs_mode"])
887 model_group.add_argument("--use-fp64-gumbel", **model_kwargs["use_fp64_gumbel"])
888 + model_group.add_argument(
889 + "--enable-trace-replay", **model_kwargs["enable_trace_replay"]
890 + )
891 model_group.add_argument(
892 "--disable-sliding-window", **model_kwargs["disable_sliding_window"]
893 )
@@ -1760,6 +1764,7 @@ class EngineArgs:
1764 max_logprobs=self.max_logprobs,
1765 logprobs_mode=self.logprobs_mode,
1766 use_fp64_gumbel=self.use_fp64_gumbel,
1767 + enable_trace_replay=self.enable_trace_replay,
1768 disable_sliding_window=self.disable_sliding_window,
1769 disable_cascade_attn=self.disable_cascade_attn,
1770 skip_tokenizer_init=self.skip_tokenizer_init,
vllm/sampling_params.py
+72 -8
@@ -346,14 +346,6 @@ class SamplingParams(
346 """Arbitrary additional args, that can be used by custom sampling
347 implementations, plugins, etc. Not used by any in-tree sampling
348 implementations."""
349 - routed_experts_prompt_start: int = 0
350 - """When enable_return_routed_experts is active, skip the first
351 - routed_experts_prompt_start prompt tokens from the returned routing
352 - data. In multi-turn agent scenarios, set this to the length of the
353 - already-returned prefix to avoid duplicating routing for prompt tokens
354 - covered by earlier turns. Default 0 returns routing for all prompt
355 - tokens."""
356 -
349 # Fields used for bad words
350 bad_words: list[str] | None = None
351 """Words that are not allowed to be generated. More precisely, only the
@@ -373,6 +365,19 @@ class SamplingParams(
365 '\\emoji \\emoji \\emoji ...'). This feature can detect such behavior
366 and terminate early, saving time and tokens."""
367
368 + # Debugging / RL-specific parameters. Not intended for production serving.
369 + routed_experts_prompt_start: int = 0
370 + """When enable_return_routed_experts is active, skip the first
371 + routed_experts_prompt_start prompt tokens from the returned routing
372 + data. In multi-turn agent scenarios, set this to the length of the
373 + already-returned prefix to avoid duplicating routing for prompt tokens
374 + covered by earlier turns. Default 0 returns routing for all prompt
375 + tokens."""
376 + trace_decode_token_ids: list[int] | None = None
377 + """If provided, forces the engine to emit this predetermined sequence of
378 + token IDs during decoding instead of sampling randomly. Real logprobs are
379 + still computed. Conflict checking is performed at the engine level."""
380 +
381 @staticmethod
382 def from_optional(
383 n: int | None = 1,
@@ -407,6 +412,8 @@ class SamplingParams(
412 repetition_detection: RepetitionDetectionParams | None = None,
413 logprob_token_ids: list[int] | None = None,
414 routed_experts_prompt_start: int = 0,
415 + # Debugging / RL-specific parameters.
416 + trace_decode_token_ids: list[int] | None = None,
417 ) -> "SamplingParams":
418 if logit_bias is not None:
419 # Fast path uses a dict comprehension; on failure we iterate once
@@ -470,6 +477,7 @@ class SamplingParams(
477 skip_clone=skip_clone,
478 repetition_detection=repetition_detection,
479 routed_experts_prompt_start=routed_experts_prompt_start,
480 + trace_decode_token_ids=trace_decode_token_ids,
481 )
482
483 def __post_init__(self) -> None:
@@ -780,6 +788,7 @@ class SamplingParams(
788 ) -> None:
789 self._validate_logprobs(model_config)
790 self._validate_logit_bias(model_config)
791 + self._validate_trace_replay(model_config, speculative_config)
792 self._validate_logits_processors(model_config)
793 self._validate_allowed_token_ids(tokenizer)
794 self._validate_spec_decode(speculative_config)
@@ -870,6 +879,61 @@ class SamplingParams(
879 value=invalid_token_ids,
880 )
881
882 + def _validate_trace_replay(
883 + self,
884 + model_config: ModelConfig,
885 + speculative_config: SpeculativeConfig | None,
886 + ) -> None:
887 + """Validate trace replay request compatibility."""
888 + if self.trace_decode_token_ids is None:
889 + return
890 +
891 + if len(self.trace_decode_token_ids) == 0:
892 + raise ValueError("trace_decode_token_ids must be a non-empty list.")
893 + if self.n != 1:
894 + raise ValueError("trace_decode_token_ids requires n=1.")
895 + if not all(isinstance(t, int) and t >= 0 for t in self.trace_decode_token_ids):
896 + raise ValueError(
897 + "trace_decode_token_ids must contain non-negative integers."
898 + )
899 +
900 + if self.prompt_logprobs is not None:
901 + raise ValueError(
902 + "trace_decode_token_ids is not supported with prompt_logprobs."
903 + )
904 + if speculative_config is not None:
905 + raise ValueError(
906 + "trace_decode_token_ids is not supported with speculative decoding."
907 + )
908 + if self.structured_outputs is not None:
909 + raise ValueError(
910 + "trace_decode_token_ids is not supported with structured outputs."
911 + )
912 + if self.repetition_detection is not None:
913 + raise ValueError(
914 + "trace_decode_token_ids is not supported with repetition_detection."
915 + )
916 + if self.thinking_token_budget is not None:
917 + raise ValueError(
918 + "trace_decode_token_ids is not supported with thinking_token_budget."
919 + )
920 + if self.bad_words:
921 + raise ValueError("trace_decode_token_ids is not supported with bad_words.")
922 +
923 + vocab_size = model_config.get_vocab_size()
924 + invalid_token_ids = [
925 + token_id
926 + for token_id in self.trace_decode_token_ids
927 + if token_id < 0 or token_id >= vocab_size
928 + ]
929 + if invalid_token_ids:
930 + raise VLLMValidationError(
931 + f"token_id(s) {invalid_token_ids} in trace_decode_token_ids "
932 + f"contain out-of-vocab token ids. Vocabulary size: {vocab_size}",
933 + parameter="trace_decode_token_ids",
934 + value=invalid_token_ids,
935 + )
936 +
937 def _validate_logits_processors(self, model_config: ModelConfig) -> None:
938 from vllm.v1.sample.logits_processor import (
939 validate_logits_processors_parameters,
vllm/v1/engine/input_processor.py
+40 -1
@@ -51,7 +51,6 @@ class InputProcessor:
51 self.speculative_config = vllm_config.speculative_config
52 self.structured_outputs_config = vllm_config.structured_outputs_config
53 self.observability_config = vllm_config.observability_config
54 - self.use_v2_model_runner = vllm_config.use_v2_model_runner
54
55 self.generation_config_fields = model_config.try_get_generation_config()
56
@@ -128,6 +127,15 @@ class InputProcessor:
127 "not configured. Please set --reasoning-parser "
128 "and/or --reasoning-config to use thinking_token_budget."
129 )
130 + if (
131 + params.trace_decode_token_ids
132 + and not self.model_config.enable_trace_replay
133 + ):
134 + raise VLLMValidationError(
135 + "trace_decode_token_ids is set but trace replay is not "
136 + "enabled. Start the engine with --enable-trace-replay "
137 + "to use it."
138 + )
139 elif isinstance(params, PoolingParams):
140 supported_pooling_tasks = [
141 task for task in supported_tasks if task in POOLING_TASKS
@@ -156,6 +164,30 @@ class InputProcessor:
164 f"but got {type(params).__name__}"
165 )
166
167 + def _normalize_trace_replay_params(
168 + self, sampling_params: SamplingParams, prompt_len: int
169 + ) -> None:
170 + """Apply trace replay's generation semantics to request-local params."""
171 + trace_token_ids = sampling_params.trace_decode_token_ids
172 + assert trace_token_ids
173 + assert sampling_params.max_tokens is not None
174 +
175 + max_trace_len = max(self.model_config.max_model_len - prompt_len, 1)
176 + trace_token_ids = trace_token_ids[:max_trace_len]
177 + sampling_params.trace_decode_token_ids = trace_token_ids
178 +
179 + # Apply this after the generation config so its EOS token cannot stop
180 + # replay before the trace is exhausted.
181 + sampling_params.max_tokens = min(
182 + len(trace_token_ids), sampling_params.max_tokens
183 + )
184 + sampling_params.min_tokens = 0
185 + sampling_params.ignore_eos = True
186 + sampling_params._eos_token_id = None
187 + sampling_params.stop = []
188 + sampling_params.stop_token_ids = []
189 + sampling_params._all_stop_token_ids = set()
190 +
191 def _validate_lora(self, lora_request: LoRARequest | None) -> None:
192 if lora_request is None:
193 return
@@ -340,6 +372,13 @@ class InputProcessor:
372 )
373 if self.tokenizer is not None:
374 sampling_params.update_from_tokenizer(self.tokenizer)
375 + if sampling_params.trace_decode_token_ids:
376 + self._normalize_trace_replay_params(
377 + sampling_params,
378 + length_from_prompt_token_ids_or_embeds(
379 + prompt_token_ids, prompt_embeds
380 + ),
381 + )
382 else:
383 pooling_params = params.clone()
384
vllm/v1/worker/gpu/model_runner.py
+1
@@ -408,6 +408,7 @@ class GPUModelRunner(LoRAModelRunnerMixin):
408 logprobs_mode=self.model_config.logprobs_mode,
409 num_speculative_tokens=self.decode_query_len,
410 use_fp64_gumbel=self.model_config.use_fp64_gumbel,
411 + enable_trace_replay=self.model_config.enable_trace_replay,
412 reasoning_config=self.vllm_config.reasoning_config,
413 return_sampling_mask=self.model_config.return_sampling_mask,
414 )
vllm/v1/worker/gpu/sample/sampler.py
+14
@@ -26,6 +26,7 @@ from vllm.v1.worker.gpu.sample.output import SamplerOutput, SamplingMaskTensors
26 from vllm.v1.worker.gpu.sample.penalties import PenaltiesState
27 from vllm.v1.worker.gpu.sample.states import NO_LOGPROBS, SamplingStates
28 from vllm.v1.worker.gpu.sample.thinking_budget import ThinkingBudgetState
29 +from vllm.v1.worker.gpu.sample.trace_replay import TraceReplayState
30 from vllm.v1.worker.gpu.states import RequestState
31
32
@@ -39,6 +40,7 @@ class Sampler:
40 logprobs_mode: LogprobsMode = "raw_logprobs",
41 num_speculative_tokens: int = 1,
42 use_fp64_gumbel: bool = False,
43 + enable_trace_replay: bool = False,
44 reasoning_config: ReasoningConfig | None = None,
45 return_sampling_mask: bool = False,
46 ):
@@ -53,6 +55,9 @@ class Sampler:
55 self.bad_words_state = BadWordsState(req_states)
56 self.logprob_token_ids_state = LogprobTokenIdsState(max_num_reqs, device)
57 self.thinking_budget_state = ThinkingBudgetState(req_states, reasoning_config)
58 + self.trace_replay_state = (
59 + TraceReplayState(req_states) if enable_trace_replay else None
60 + )
61 self.needs_logits_processing = np.zeros(max_num_reqs, dtype=bool)
62 self.num_speculative_tokens = num_speculative_tokens
63 self.return_sampling_mask = return_sampling_mask
@@ -69,6 +74,8 @@ class Sampler:
74 self.bad_words_state.add_request(req_idx, sampling_params)
75 self.logprob_token_ids_state.add_request(req_idx, sampling_params)
76 self.thinking_budget_state.add_request(req_idx, sampling_params)
77 + if self.trace_replay_state is not None:
78 + self.trace_replay_state.add_request(req_idx, sampling_params)
79
80 states = self.sampling_states
81 temperature = states.temperature.np[req_idx]
@@ -93,6 +100,8 @@ class Sampler:
100 self.bad_words_state.apply_staged_writes()
101 self.logprob_token_ids_state.apply_staged_writes()
102 self.thinking_budget_state.apply_staged_writes()
103 + if self.trace_replay_state is not None:
104 + self.trace_replay_state.apply_staged_writes()
105
106 def __call__(
107 self,
@@ -128,6 +137,11 @@ class Sampler:
137 return_logprobs=return_logprobs,
138 )
139
140 + if self.trace_replay_state is not None:
141 + # Overwrite sampled tokens with the replay trace up-front so that
142 + # computed logprobs reflect the real distribution of the forced token.
143 + self.trace_replay_state.apply_trace(sampled, idx_mapping)
144 +
145 if return_logprobs:
146 if self.logprobs_mode in PROCESSED_LOGPROBS_MODES:
147 logits = processed_logits
vllm/v1/worker/gpu/sample/trace_replay.py new
+113
@@ -0,0 +1,113 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +import torch
4 +
5 +from vllm.sampling_params import SamplingParams
6 +from vllm.triton_utils import tl, triton
7 +from vllm.v1.worker.gpu.buffer_utils import StagedWriteTensor, UvaBackedTensor
8 +from vllm.v1.worker.gpu.states import RequestState
9 +
10 +
11 +class TraceReplayState:
12 + """Per-request state for inference trace-replay.
13 +
14 + When a request carries ``SamplingParams.trace_decode_token_ids``, the
15 + sampler overwrites the sampled token at each decode step with the
16 + predetermined trace token, while real logprobs and ranks are still computed
17 + from the unmodified logit distribution. The replay step for a request is
18 + derived entirely from GPU state (``total_len - prompt_len``), so no CPU
19 + synchronization or async placeholder handling is needed.
20 + """
21 +
22 + def __init__(self, req_states: RequestState):
23 + self.max_num_reqs = req_states.max_num_reqs
24 + self.device = req_states.device
25 + self.req_states = req_states
26 + self.trace_token_ids = StagedWriteTensor(
27 + (self.max_num_reqs, req_states.max_model_len),
28 + dtype=torch.int32,
29 + device=self.device,
30 + uva_instead_of_gpu=True,
31 + )
32 + self.trace_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)
33 +
34 + def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
35 + trace = sampling_params.trace_decode_token_ids
36 + if trace is not None:
37 + self.trace_len.np[req_idx] = len(trace)
38 + self.trace_token_ids.stage_write(req_idx, 0, trace)
39 + else:
40 + self.trace_len.np[req_idx] = 0
41 +
42 + def apply_staged_writes(self) -> None:
43 + self.trace_len.copy_to_uva()
44 + self.trace_token_ids.apply_write()
45 +
46 + def apply_trace(
47 + self,
48 + sampled: torch.Tensor,
49 + idx_mapping: torch.Tensor,
50 + ) -> None:
51 + apply_trace_tokens(
52 + sampled,
53 + idx_mapping,
54 + self.trace_token_ids.gpu,
55 + self.trace_len.gpu,
56 + self.req_states.total_len.gpu,
57 + self.req_states.prompt_len.gpu,
58 + )
59 +
60 +
61 +@triton.jit
62 +def _trace_replay_kernel(
63 + sampled_ptr, # [num_reqs], int64, mutated in place
64 + idx_mapping_ptr, # [num_reqs] batch_idx -> req_state_idx
65 + trace_token_ids_ptr, # [max_num_reqs, max_model_len], int32
66 + trace_token_ids_stride,
67 + trace_len_ptr, # [max_num_reqs], int32
68 + total_len_ptr, # [max_num_reqs], int32
69 + prompt_len_ptr, # [max_num_reqs], int32
70 +):
71 + batch_idx = tl.program_id(0)
72 + req_state_idx = tl.load(idx_mapping_ptr + batch_idx)
73 + if req_state_idx < 0:
74 + return
75 +
76 + trace_len = tl.load(trace_len_ptr + req_state_idx)
77 + if trace_len <= 0:
78 + return
79 +
80 + # The token being sampled now is output token number
81 + # (total_len - prompt_len): total_len reflects tokens committed through the
82 + # previous step (post_update runs after sampling).
83 + step = tl.load(total_len_ptr + req_state_idx) - tl.load(
84 + prompt_len_ptr + req_state_idx
85 + )
86 + if step < 0 or step >= trace_len:
87 + return
88 +
89 + token_id = tl.load(
90 + trace_token_ids_ptr + req_state_idx * trace_token_ids_stride + step
91 + )
92 + tl.store(sampled_ptr + batch_idx, token_id.to(tl.int64))
93 +
94 +
95 +def apply_trace_tokens(
96 + sampled: torch.Tensor,
97 + idx_mapping: torch.Tensor,
98 + trace_token_ids: torch.Tensor,
99 + trace_len: torch.Tensor,
100 + total_len: torch.Tensor,
101 + prompt_len: torch.Tensor,
102 +) -> None:
103 + """Overwrite ``sampled`` in place with trace tokens for the current step."""
104 + num_reqs = idx_mapping.shape[0]
105 + _trace_replay_kernel[(num_reqs,)](
106 + sampled,
107 + idx_mapping,
108 + trace_token_ids,
109 + trace_token_ids.stride(0),
110 + trace_len,
111 + total_len,
112 + prompt_len,
113 + )