vllm.v1.worker.gpu.sample.trace_replay

Classes:

• TraceReplayState –

  Per-request state for inference trace-replay.

Functions:

• apply_trace_tokens –

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

• apply_trace_tokens_to_verified –

  Replace verified tokens with the trace, accepting each draft only if it

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

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

    def apply_trace_to_verified(
        self,
        sampled: torch.Tensor,
        num_sampled: torch.Tensor,
        draft_sampled: torch.Tensor,
        cu_num_logits: torch.Tensor,
        idx_mapping: torch.Tensor,
    ) -> None:
        apply_trace_tokens_to_verified(
            sampled,
            num_sampled,
            draft_sampled,
            cu_num_logits,
            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

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

apply_trace_tokens_to_verified(sampled, num_sampled, draft_sampled, cu_num_logits, idx_mapping, trace_token_ids, trace_len, total_len, prompt_len)

Replace verified tokens with the trace, accepting each draft only if it equals the trace token at its position.

Source code in vllm/v1/worker/gpu/sample/trace_replay.py

def apply_trace_tokens_to_verified(
    sampled: torch.Tensor,
    num_sampled: torch.Tensor,
    draft_sampled: torch.Tensor,
    cu_num_logits: torch.Tensor,
    idx_mapping: torch.Tensor,
    trace_token_ids: torch.Tensor,
    trace_len: torch.Tensor,
    total_len: torch.Tensor,
    prompt_len: torch.Tensor,
) -> None:
    """Replace verified tokens with the trace, accepting each draft only if it
    equals the trace token at its position."""
    num_reqs = idx_mapping.shape[0]
    _trace_replay_verified_kernel[(num_reqs,)](
        sampled,
        sampled.stride(0),
        num_sampled,
        draft_sampled,
        cu_num_logits,
        idx_mapping,
        trace_token_ids,
        trace_token_ids.stride(0),
        trace_len,
        total_len,
        prompt_len,
    )