| 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] |