vllm.v1.worker.gpu.async_utils

Classes:

• StepTimingCollector –

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

Functions:

• stream –

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

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)