@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 141 lines 5.03 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 from collections.abc import Callable
5 from enum import Enum
6 from typing import Any
7
8 import torch
9
10
11 class EventType(Enum):
12 Main = 0
13 Attention = 1
14
15
16 def maybe_execute_in_parallel(
17 fn0: Callable[[], Any],
18 fn1: Callable[[], Any],
19 event0: torch.cuda.Event,
20 event1: torch.cuda.Event,
21 aux_stream: torch.cuda.Stream | None = None,
22 ) -> tuple[Any, Any]:
23 """Run two functions potentially in parallel on separate CUDA streams.
24
25 When aux_stream is provided, fn0 runs on the current (default) stream and
26 fn1 runs on aux_stream, synchronized via CUDA events. When aux_stream is
27 None or a breakable CUDA graph capture is active, both functions execute
28 sequentially on the current stream.
29
30 This design follows TensorRT-LLM's maybe_execute_in_parallel pattern
31 (tensorrt_llm/_torch/modules/multi_stream_utils.py).
32
33 Args:
34 fn0: Callable for the default stream.
35 fn1: Callable for the auxiliary stream.
36 event0: CUDA event recorded before fn0 so aux_stream can wait.
37 event1: CUDA event recorded after fn1 so default stream can wait.
38 aux_stream: The second CUDA stream for fn1.
39 Multi-stream is disabled when aux_stream is None or a breakable
40 CUDA graph capture is active.
41
42 Returns:
43 Tuple of (fn0_result, fn1_result).
44
45 """
46 if aux_stream is not None:
47 from vllm.compilation.breakable_cudagraph import BreakableCUDAGraphCapture
48
49 if BreakableCUDAGraphCapture.is_active():
50 aux_stream = None
51
52 if aux_stream is not None:
53 event0.record()
54 result0 = fn0()
55 with torch.cuda.stream(aux_stream):
56 event0.wait()
57 result1 = fn1()
58 event1.record()
59 event1.wait()
60 else:
61 result0 = fn0()
62 result1 = fn1()
63 return (result0, result1)
64
65
66 def execute_in_parallel(
67 default_fn: Callable[[], Any],
68 aux_fns: list[Callable[[], Any] | None],
69 start_event: torch.cuda.Event,
70 done_events: list[torch.cuda.Event],
71 aux_streams: list[torch.cuda.Stream] | None = None,
72 enable: bool = False,
73 default_first: bool = False,
74 ) -> tuple[Any, list[Any]]:
75 """Run default_fn on the current stream and aux_fns concurrently on
76 aux_streams.
77
78 Generalizes maybe_execute_in_parallel to N aux callables. Slots where
79 aux_fns[i] is None are skipped (no stream switch, no event record); their
80 corresponding entry in the returned aux_results list is None.
81
82 start_event fans out from the current stream to every launched aux stream;
83 done_events[i] is recorded after aux_fns[i] so the current stream joins
84 before returning. Falls back to sequential execution on the current stream
85 when aux_streams is None or enable is False; in that case default_fn runs
86 first, then aux_fns in order.
87
88 Args:
89 default_fn: Callable for the default (current) stream.
90 aux_fns: Per-aux callables; entries may be None to skip.
91 start_event: CUDA event recorded on the current stream before
92 default_fn so each launched aux stream can wait on it.
93 done_events: One CUDA event per aux slot, recorded after the
94 corresponding aux_fn. Length must match aux_fns.
95 aux_streams: Per-aux CUDA streams. Length must match aux_fns.
96 Multi-stream is disabled when None.
97 enable: Opt-in switch for the multi-stream path. Defaults to False,
98 so callers that pass aux_streams must also pass enable=True
99 (typically gated by an env var) to actually overlap. When False,
100 execution falls back to sequential on the current stream.
101 default_first: When True, enqueue the default chain before the aux
102 chains. The default chain is typically the longest, so it starts
103 first; the aux streams still only wait on start_event, so they
104 overlap the default chain. Defaults to False.
105
106 Returns:
107 Tuple of (default_result, aux_results) where aux_results[i] is the
108 result of aux_fns[i] (or None when skipped).
109
110 """
111 aux_results: list[Any]
112 if aux_streams is None or not enable:
113 default_result = default_fn()
114 aux_results = [fn() if fn is not None else None for fn in aux_fns]
115 return default_result, aux_results
116
117 assert len(aux_fns) == len(aux_streams) == len(done_events), (
118 "aux_fns, aux_streams, and done_events must be the same length"
119 )
120
121 aux_results = [None] * len(aux_fns)
122 pending: list[torch.cuda.Event] = []
123
124 start_event.record()
125 if default_first:
126 default_result = default_fn()
127 for i, fn in enumerate(aux_fns):
128 if fn is None:
129 continue
130 with torch.cuda.stream(aux_streams[i]):
131 start_event.wait()
132 aux_results[i] = fn()
133 done_events[i].record()
134 pending.append(done_events[i])
135 if not default_first:
136 default_result = default_fn()
137
138 for ev in pending:
139 ev.wait()
140
141 return default_result, aux_results