Skip to content

vllm.v1.worker.gpu.sample.trace_replay

Classes:

Functions:

  • apply_trace_tokens –

    Overwrite sampled in place with trace tokens for the current step.

TraceReplayState

Per-request state for inference trace-replay.

When a request carries SamplingParams.trace_decode_token_ids, the sampler overwrites the sampled token at each decode step with the predetermined trace token, while real logprobs and ranks are still computed from the unmodified logit distribution. The replay step for a request is derived entirely from GPU state (total_len - prompt_len), so no CPU synchronization or async placeholder handling is needed.

Source code in vllm/v1/worker/gpu/sample/trace_replay.py 
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
class TraceReplayState:
    """Per-request state for inference trace-replay.

    When a request carries ``SamplingParams.trace_decode_token_ids``, the
    sampler overwrites the sampled token at each decode step with the
    predetermined trace token, while real logprobs and ranks are still computed
    from the unmodified logit distribution. The replay step for a request is
    derived entirely from GPU state (``total_len - prompt_len``), so no CPU
    synchronization or async placeholder handling is needed.
    """

    def __init__(self, req_states: RequestState):
        self.max_num_reqs = req_states.max_num_reqs
        self.device = req_states.device
        self.req_states = req_states
        self.trace_token_ids = StagedWriteTensor(
            (self.max_num_reqs, req_states.max_model_len),
            dtype=torch.int32,
            device=self.device,
            uva_instead_of_gpu=True,
        )
        self.trace_len = UvaBackedTensor(self.max_num_reqs, dtype=torch.int32)

    def add_request(self, req_idx: int, sampling_params: SamplingParams) -> None:
        trace = sampling_params.trace_decode_token_ids
        if trace is not None:
            self.trace_len.np[req_idx] = len(trace)
            self.trace_token_ids.stage_write(req_idx, 0, trace)
        else:
            self.trace_len.np[req_idx] = 0

    def apply_staged_writes(self) -> None:
        self.trace_len.copy_to_uva()
        self.trace_token_ids.apply_write()

    def apply_trace(
        self,
        sampled: torch.Tensor,
        idx_mapping: torch.Tensor,
    ) -> None:
        apply_trace_tokens(
            sampled,
            idx_mapping,
            self.trace_token_ids.gpu,
            self.trace_len.gpu,
            self.req_states.total_len.gpu,
            self.req_states.prompt_len.gpu,
        )

apply_trace_tokens(sampled, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len)

Overwrite sampled in place with trace tokens for the current step.

Source code in vllm/v1/worker/gpu/sample/trace_replay.py 
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
def apply_trace_tokens(
    sampled: torch.Tensor,
    idx_mapping: torch.Tensor,
    trace_token_ids: torch.Tensor,
    trace_len: torch.Tensor,
    total_len: torch.Tensor,
    prompt_len: torch.Tensor,
) -> None:
    """Overwrite ``sampled`` in place with trace tokens for the current step."""
    num_reqs = idx_mapping.shape[0]
    _trace_replay_kernel[(num_reqs,)](
        sampled,
        idx_mapping,
        trace_token_ids,
        trace_token_ids.stride(0),
        trace_len,
        total_len,
        prompt_len,
    )