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