Skip to content

vllm.v1.worker.gpu.async_utils

Classes:

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 
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
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 
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
@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 
107
108
109
110
111
112
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 
90
91
92
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 
257
258
259
260
261
262
263
264
265
266
@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)