vllm.v1.worker.gpu.async_utils

Classes:

• AsyncOutput –
• AsyncPoolingOutput –
• StepTimingCollector –

  Times the steps run inside collect; record calls no-op outside it.

Functions:

• stream –

  Lightweight version of torch.cuda.stream() context manager which

AsyncOutput

Bases: AsyncModelRunnerOutput

Methods:

• get_kv_connector_output –

  Get the KV connector output for this async output.

Source code in vllm/v1/worker/gpu/async_utils.py

class AsyncOutput(AsyncModelRunnerOutput):
    def __init__(
        self,
        model_runner_output: ModelRunnerOutput,
        sampler_output: SamplerOutput,
        num_sampled_tokens: torch.Tensor,
        main_stream: torch.cuda.Stream,
        copy_stream: torch.cuda.Stream,
        check_ep_fault: bool,
        pending_aux_output: "PendingAuxOutput | None",
    ):
        # NOTE(woosuk): We must retain references to the GPU tensors,
        # as the copy operations are performed on a different CUDA stream than
        # the one where the tensors were created.
        self.model_runner_output = model_runner_output
        self.sampler_output = sampler_output
        self.num_sampled_tokens = num_sampled_tokens
        self.pending_aux_output = pending_aux_output
        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
        self.copy_event = torch.cuda.Event(blocking=True)
        self._has_fault: torch.Tensor | None = None

        with stream(copy_stream, main_stream):
            copy_stream.wait_stream(main_stream)

            self.sampled_token_ids = async_copy_to_np(sampler_output.sampled_token_ids)
            self.logprobs_tensors: LogprobsTensors | None = None
            if sampler_output.logprobs_tensors is not None:
                self.logprobs_tensors = (
                    sampler_output.logprobs_tensors.to_cpu_nonblocking()
                )
            self.num_nans: np.ndarray | None = None
            if sampler_output.num_nans is not None:
                self.num_nans = async_copy_to_np(sampler_output.num_nans)
            self.num_sampled_tokens_np = async_copy_to_np(num_sampled_tokens)
            self.sampling_mask_tensors: SamplingMaskTensors | None = None
            if sampler_output.sampling_mask_tensors is not None:
                self.sampling_mask_tensors = (
                    sampler_output.sampling_mask_tensors.to_cpu_nonblocking()
                )
            self.prompt_logprobs_dict = {
                k: v.to_cpu_nonblocking() if v is not None else None
                for k, v in self.model_runner_output.prompt_logprobs_dict.items()
            }
            token_id_logprobs = self.model_runner_output.prompt_token_id_logprobs_dict
            self.prompt_token_id_logprobs_dict = {
                k: v.to("cpu", non_blocking=True) for k, v in token_id_logprobs.items()
            }
            if self.pending_aux_output is not None:
                self.pending_aux_output.enqueue_cpu_copy(
                    num_sampled=self.num_sampled_tokens_np,
                    num_rejected=async_copy_to_np(sampler_output.num_rejected),
                )
            if check_ep_fault:
                has_fault = get_ep_all2all_manager().query_fault()
                self._has_fault = has_fault.to("cpu", non_blocking=True)
            self.copy_event.record(copy_stream)

    def get_output(self) -> ModelRunnerOutput:
        self.copy_event.synchronize()

        # NOTE(woosuk): The following code is to ensure compatibility with
        # the existing model runner.
        # Going forward, we should keep the data structures as NumPy arrays
        # rather than Python lists.
        sampled_token_ids: list[list[int]] = self.sampled_token_ids.tolist()
        num_sampled_tokens: list[int] = self.num_sampled_tokens_np.tolist()
        for token_ids, num_tokens in zip(sampled_token_ids, num_sampled_tokens):
            del token_ids[num_tokens:]
        self.model_runner_output.sampled_token_ids = sampled_token_ids

        if self.sampling_mask_tensors is not None:
            self.model_runner_output.sampling_masks = (
                self.sampling_mask_tensors.tolists(self.num_sampled_tokens_np)
            )

        if self.num_nans is not None:
            self.model_runner_output.num_nans_in_logits = dict(
                zip(self.model_runner_output.req_ids, self.num_nans.tolist())
            )
            if envs.VLLM_RAISE_ON_LOGIT_NANS:
                raise_if_nan_logits(self.model_runner_output.num_nans_in_logits)

        if self.logprobs_tensors is not None:
            self.model_runner_output.logprobs = self.logprobs_tensors.tolists()
        self.model_runner_output.prompt_logprobs_dict = self.prompt_logprobs_dict
        self.model_runner_output.prompt_token_id_logprobs_dict = (
            self.prompt_token_id_logprobs_dict
        )
        if self.pending_aux_output is not None:
            self.model_runner_output.aux_output_connector_output = (
                self.pending_aux_output.process_output()
            )

        if self._has_fault is not None and self._has_fault.item():
            mask = get_ep_all2all_manager().query_active_mask()
            raise RuntimeError(
                "Fault detected in EP all2all communication: "
                "one or more ranks timed out during dispatch/combine. "
                f"Mask: {mask.cpu().tolist()}"
            )

        return self.model_runner_output

    def get_kv_connector_output(self) -> KVConnectorOutput | None:
        """Get the KV connector output for this async output."""
        return self.model_runner_output.kv_connector_output

get_kv_connector_output()

Get the KV connector output for this async output.

Source code in vllm/v1/worker/gpu/async_utils.py

def get_kv_connector_output(self) -> KVConnectorOutput | None:
    """Get the KV connector output for this async output."""
    return self.model_runner_output.kv_connector_output

AsyncPoolingOutput

Bases: AsyncModelRunnerOutput

Methods:

• get_kv_connector_output –

  Get the KV connector output for this async output.

Source code in vllm/v1/worker/gpu/async_utils.py

class AsyncPoolingOutput(AsyncModelRunnerOutput):
    def __init__(
        self,
        model_runner_output: ModelRunnerOutput,
        pooler_output: PoolerOutput,
        finished_mask: list[bool],
        main_stream: torch.cuda.Stream,
        copy_stream: torch.cuda.Stream,
    ):
        self.model_runner_output = model_runner_output
        self.pooler_output = pooler_output
        # Blocking (sleep) event to avoid busy-polling the CUDA driver lock.
        self.copy_event = torch.cuda.Event(blocking=True)

        with stream(copy_stream, main_stream):
            copy_stream.wait_stream(main_stream)
            if isinstance(self.pooler_output, torch.Tensor) and all(finished_mask):
                self.pooler_output_cpu: PoolerOutput = self.pooler_output.to(
                    "cpu", non_blocking=True
                )
            else:
                outputs = (
                    self.pooler_output.unbind()
                    if isinstance(self.pooler_output, torch.Tensor)
                    else self.pooler_output
                )
                self.pooler_output_cpu = [
                    None
                    if output is None or not is_finished
                    else output.to("cpu", non_blocking=True)
                    for output, is_finished in zip(outputs, finished_mask, strict=True)
                ]
            self.copy_event.record(copy_stream)

    def get_output(self) -> ModelRunnerOutput:
        if isinstance(self.pooler_output_cpu, torch.Tensor):
            pooler_output = list(self.pooler_output_cpu.unbind(dim=0))
        else:
            pooler_output = self.pooler_output_cpu
        self.copy_event.synchronize()
        self.model_runner_output.pooler_output = pooler_output
        return self.model_runner_output

    def get_kv_connector_output(self) -> KVConnectorOutput | None:
        """Get the KV connector output for this async output."""
        return self.model_runner_output.kv_connector_output

get_kv_connector_output()

Get the KV connector output for this async output.

Source code in vllm/v1/worker/gpu/async_utils.py

def get_kv_connector_output(self) -> KVConnectorOutput | None:
    """Get the KV connector output for this async output."""
    return self.model_runner_output.kv_connector_output

StepTimingCollector

Times the steps run inside collect; record calls no-op outside it.

Every step gets its own events, so the steps queue back-to-back and the block resolves them all behind one sync on the way out.

Methods:

• collect –

  Time every step run in this block.

• drafter_end –

  Ends the step: only steps that reach here have a draft cost.

• record_batch –

  Costs from different execution modes must not share a cost curve.

Source code in vllm/v1/worker/gpu/async_utils.py

class StepTimingCollector:
    """Times the steps run inside ``collect``; record calls no-op outside it.

    Every step gets its own events, so the steps queue back-to-back and the
    block resolves them all behind one sync on the way out.
    """

    def __init__(self):
        self._collecting = False
        self._step: StepTimingEvents | None = None
        self._batch = (False, 0, 0)
        self._timed: list[tuple[StepTimingEvents, tuple[bool, int, int]]] = []

    @contextlib.contextmanager
    def collect(self) -> Iterator[list[StepTimingSample]]:
        """Time every step run in this block.

        The yielded list holds one sample per timed step once the block exits;
        it stays empty inside the block, where the timings are still on device.
        """
        samples: list[StepTimingSample] = []
        self._collecting = True
        try:
            yield samples
        finally:
            self._collecting = False
            timed, self._timed, self._step = self._timed, [], None
        if not timed:
            return
        # Same stream, issue order: once the last step is done, so are the rest.
        timed[-1][0].drafter_end.synchronize()
        samples.extend(
            StepTimingSample(
                events.forward_start.elapsed_time(events.forward_end),
                events.drafter_start.elapsed_time(events.drafter_end),
                num_target_tokens,
                num_reqs,
                full_cudagraph,
            )
            for events, (full_cudagraph, num_target_tokens, num_reqs) in timed
        )

    def record_batch(self, input_batch: "InputBatch", full_cudagraph: bool) -> None:
        """Costs from different execution modes must not share a cost curve."""
        self._batch = (full_cudagraph, input_batch.num_tokens, input_batch.num_reqs)

    def forward_start(self) -> None:
        if self._collecting:
            self._step = StepTimingEvents()
            self._step.forward_start.record()

    def forward_end(self) -> None:
        if self._step is not None:
            self._step.forward_end.record()

    def drafter_start(self) -> None:
        if self._step is not None:
            self._step.drafter_start.record()

    def drafter_end(self) -> None:
        """Ends the step: only steps that reach here have a draft cost."""
        if self._step is not None:
            self._step.drafter_end.record()
            self._timed.append((self._step, self._batch))
            self._step = None

collect()

Time every step run in this block.

The yielded list holds one sample per timed step once the block exits; it stays empty inside the block, where the timings are still on device.

Source code in vllm/v1/worker/gpu/async_utils.py

@contextlib.contextmanager
def collect(self) -> Iterator[list[StepTimingSample]]:
    """Time every step run in this block.

    The yielded list holds one sample per timed step once the block exits;
    it stays empty inside the block, where the timings are still on device.
    """
    samples: list[StepTimingSample] = []
    self._collecting = True
    try:
        yield samples
    finally:
        self._collecting = False
        timed, self._timed, self._step = self._timed, [], None
    if not timed:
        return
    # Same stream, issue order: once the last step is done, so are the rest.
    timed[-1][0].drafter_end.synchronize()
    samples.extend(
        StepTimingSample(
            events.forward_start.elapsed_time(events.forward_end),
            events.drafter_start.elapsed_time(events.drafter_end),
            num_target_tokens,
            num_reqs,
            full_cudagraph,
        )
        for events, (full_cudagraph, num_target_tokens, num_reqs) in timed
    )

drafter_end()

Ends the step: only steps that reach here have a draft cost.

Source code in vllm/v1/worker/gpu/async_utils.py

def drafter_end(self) -> None:
    """Ends the step: only steps that reach here have a draft cost."""
    if self._step is not None:
        self._step.drafter_end.record()
        self._timed.append((self._step, self._batch))
        self._step = None

record_batch(input_batch, full_cudagraph)

Costs from different execution modes must not share a cost curve.

Source code in vllm/v1/worker/gpu/async_utils.py

def record_batch(self, input_batch: "InputBatch", full_cudagraph: bool) -> None:
    """Costs from different execution modes must not share a cost curve."""
    self._batch = (full_cudagraph, input_batch.num_tokens, input_batch.num_reqs)

stream(to_stream, from_stream)

Lightweight version of torch.cuda.stream() context manager which avoids current_stream and device lookups.

Source code in vllm/v1/worker/gpu/async_utils.py

@contextlib.contextmanager
def stream(to_stream: torch.cuda.Stream, from_stream: torch.cuda.Stream):
    """Lightweight version of torch.cuda.stream() context manager which
    avoids current_stream and device lookups.
    """
    try:
        torch.cuda.set_stream(to_stream)
        yield
    finally:
        torch.cuda.set_stream(from_stream)