|
| 1 | +# SPDX-License-Identifier: Apache-2.0 |
| 2 | +# DeepSpeed Team |
| 3 | +"""CUDA graph capture for HybridEngine token generation. |
| 4 | +
|
| 5 | +Generation in the HybridEngine is bound by CPU work rather than by the GPU: a |
| 6 | +single decode step issues on the order of a thousand kernel launches, and the |
| 7 | +kernels finish well before the host can queue the next step. Capturing the decode |
| 8 | +forward into a CUDA graph replaces all of those launches with one replay. |
| 9 | +
|
| 10 | +Two properties of the inference kernels shape this design. |
| 11 | +
|
| 12 | +First, a single graph cannot serve every decode position. The kernels read the |
| 13 | +current sequence length from a host-side counter |
| 14 | +(``InferenceContext::current_tokens()``) and pass it to kernels as a launch |
| 15 | +parameter. Capture freezes launch parameters, so a graph captured at one position |
| 16 | +would keep reading and writing that same position. Host code *does* run during |
| 17 | +capture, so capturing one graph per position records the correct offsets into |
| 18 | +each. |
| 19 | +
|
| 20 | +Second, that counter is advanced from host code and is not reachable from Python. |
| 21 | +Replay runs no host code, so it leaves the counter behind. An eager decode step |
| 22 | +after a replay would therefore use a stale sequence length and silently corrupt |
| 23 | +the KV cache. Eager and replayed decode steps must never be mixed within one |
| 24 | +sequence, which is why the decision is made once per sequence, up front, in |
| 25 | +``begin_sequence``. |
| 26 | +""" |
| 27 | + |
| 28 | +import torch |
| 29 | + |
| 30 | +from deepspeed.accelerator import get_accelerator |
| 31 | +from deepspeed.utils import logger |
| 32 | + |
| 33 | + |
| 34 | +def _tensor_kwargs(kwargs): |
| 35 | + """Names of the keyword arguments that hold a tensor.""" |
| 36 | + return [name for name, value in kwargs.items() if torch.is_tensor(value)] |
| 37 | + |
| 38 | + |
| 39 | +class DecodeGraphCache: |
| 40 | + """Dispatches generation forwards to a per-position CUDA graph. |
| 41 | +
|
| 42 | + Args: |
| 43 | + forward: The original ``forward``, used both for capture and as the |
| 44 | + eager fallback. |
| 45 | + max_positions: Upper bound on how many decode positions may be captured, |
| 46 | + which bounds the memory the cache can consume. |
| 47 | + """ |
| 48 | + |
| 49 | + def __init__(self, forward, max_positions): |
| 50 | + self._forward = forward |
| 51 | + self._max_positions = max_positions |
| 52 | + |
| 53 | + self._graphs = {} |
| 54 | + self._static_kwargs = {} |
| 55 | + self._static_outputs = {} |
| 56 | + self._pool = None |
| 57 | + |
| 58 | + self._position = 0 |
| 59 | + self._captured_length = None |
| 60 | + self._use_graphs = False |
| 61 | + self._disabled = False |
| 62 | + |
| 63 | + @property |
| 64 | + def captured_positions(self): |
| 65 | + return len(self._graphs) |
| 66 | + |
| 67 | + def invalidate(self): |
| 68 | + """Drop every captured graph and the memory pool backing them.""" |
| 69 | + self._graphs.clear() |
| 70 | + self._static_kwargs.clear() |
| 71 | + self._static_outputs.clear() |
| 72 | + self._pool = None |
| 73 | + self._captured_length = None |
| 74 | + |
| 75 | + def begin_sequence(self, num_decode_steps): |
| 76 | + """Decide once, before any decode step, whether this sequence uses graphs. |
| 77 | +
|
| 78 | + Args: |
| 79 | + num_decode_steps: How many single-token forwards this generate call |
| 80 | + will make, or ``None`` when that is not known ahead of time. |
| 81 | +
|
| 82 | + A sequence runs entirely on graphs or entirely eagerly. Deciding here |
| 83 | + rather than per step is what keeps the host-side sequence counter |
| 84 | + consistent: a capture pass advances it on every step, and a replay pass |
| 85 | + never reads it. |
| 86 | + """ |
| 87 | + self._position = 0 |
| 88 | + |
| 89 | + if self._disabled or num_decode_steps is None or num_decode_steps <= 0: |
| 90 | + self._use_graphs = False |
| 91 | + return |
| 92 | + |
| 93 | + if num_decode_steps > self._max_positions: |
| 94 | + self._use_graphs = False |
| 95 | + return |
| 96 | + |
| 97 | + # A different generation length needs its own set of graphs, because each |
| 98 | + # graph has both its sequence offset and its attention-mask width baked in. |
| 99 | + if self._captured_length is not None and self._captured_length != num_decode_steps: |
| 100 | + logger.info(f"HybridEngine CUDA graph: generation length changed " |
| 101 | + f"{self._captured_length} -> {num_decode_steps}, recapturing.") |
| 102 | + self.invalidate() |
| 103 | + |
| 104 | + self._captured_length = num_decode_steps |
| 105 | + self._use_graphs = True |
| 106 | + |
| 107 | + def _capture(self, position, args, kwargs): |
| 108 | + """Record a graph for ``position`` and keep its static input buffers.""" |
| 109 | + static = {name: kwargs[name].clone() for name in _tensor_kwargs(kwargs)} |
| 110 | + capture_kwargs = dict(kwargs) |
| 111 | + capture_kwargs.update(static) |
| 112 | + |
| 113 | + graph = get_accelerator().create_graph() |
| 114 | + if self._pool is None: |
| 115 | + with get_accelerator().capture_to_graph(graph): |
| 116 | + output = self._forward(*args, **capture_kwargs) |
| 117 | + self._pool = graph.pool() |
| 118 | + else: |
| 119 | + with get_accelerator().capture_to_graph(graph, pool=self._pool): |
| 120 | + output = self._forward(*args, **capture_kwargs) |
| 121 | + |
| 122 | + self._graphs[position] = graph |
| 123 | + self._static_kwargs[position] = static |
| 124 | + self._static_outputs[position] = output |
| 125 | + |
| 126 | + def _replay(self, position, kwargs): |
| 127 | + for name, buffer in self._static_kwargs[position].items(): |
| 128 | + buffer.copy_(kwargs[name]) |
| 129 | + get_accelerator().replay_graph(self._graphs[position]) |
| 130 | + return self._static_outputs[position] |
| 131 | + |
| 132 | + def _shapes_match(self, position, kwargs): |
| 133 | + captured = self._static_kwargs[position] |
| 134 | + if set(captured) != set(_tensor_kwargs(kwargs)): |
| 135 | + return False |
| 136 | + return all(captured[name].shape == kwargs[name].shape for name in captured) |
| 137 | + |
| 138 | + def _fall_back(self, reason, args, kwargs): |
| 139 | + """Give up on graphs for good and run eagerly. |
| 140 | +
|
| 141 | + Only safe before any replay has happened in the current sequence, which |
| 142 | + is why this is reachable only from the capture path. |
| 143 | + """ |
| 144 | + logger.warning(f"HybridEngine CUDA graph disabled: {reason}") |
| 145 | + self.invalidate() |
| 146 | + self._disabled = True |
| 147 | + self._use_graphs = False |
| 148 | + return self._forward(*args, **kwargs) |
| 149 | + |
| 150 | + def __call__(self, *args, **kwargs): |
| 151 | + input_ids = kwargs.get("input_ids") |
| 152 | + is_decode_step = (input_ids is not None and input_ids.dim() == 2 and input_ids.shape[1] == 1) |
| 153 | + |
| 154 | + if not self._use_graphs or not is_decode_step: |
| 155 | + return self._forward(*args, **kwargs) |
| 156 | + |
| 157 | + position = self._position |
| 158 | + self._position += 1 |
| 159 | + |
| 160 | + if position >= self._max_positions: |
| 161 | + # begin_sequence() bounds the length, so this means the caller ran |
| 162 | + # longer than it declared. Replays have already happened, so eager |
| 163 | + # execution would read a stale sequence counter. |
| 164 | + raise RuntimeError(f"HybridEngine CUDA graph: generation exceeded the declared length " |
| 165 | + f"({self._max_positions} decode steps). Set hybrid_engine.max_out_tokens " |
| 166 | + f"to cover the longest generation, or disable hybrid_engine.enable_cuda_graph.") |
| 167 | + |
| 168 | + if position in self._graphs: |
| 169 | + if self._shapes_match(position, kwargs): |
| 170 | + return self._replay(position, kwargs) |
| 171 | + return self._fall_back("input shapes changed mid-sequence", args, kwargs) |
| 172 | + |
| 173 | + try: |
| 174 | + self._capture(position, args, kwargs) |
| 175 | + except Exception as err: |
| 176 | + return self._fall_back(f"capture failed at position {position}: {err}", args, kwargs) |
| 177 | + |
| 178 | + # Capture records the work without running it, so the first execution must |
| 179 | + # come from a replay. That is also what fills the KV cache for this position. |
| 180 | + return self._replay(position, kwargs) |
| 181 | + |
| 182 | + |
| 183 | +def decode_steps_from_generate_kwargs(kwargs): |
| 184 | + """How many single-token forwards a generate call will make, if knowable. |
| 185 | +
|
| 186 | + HF emits one prompt forward that produces the first new token, then one |
| 187 | + forward per remaining token. Only a pinned length is usable here: with an |
| 188 | + open-ended limit the sequence may stop early, and a sequence that runs |
| 189 | + *longer* than its captured graphs cannot fall back safely. |
| 190 | + """ |
| 191 | + max_new = kwargs.get("max_new_tokens") |
| 192 | + min_new = kwargs.get("min_new_tokens") |
| 193 | + if max_new is None or min_new is None or max_new != min_new: |
| 194 | + return None |
| 195 | + return int(max_new) - 1 |
| 196 | + |
| 197 | + |
| 198 | +def validate_cuda_graph_support(config, zero_stage): |
| 199 | + """Return the reason CUDA graphs cannot be used, or ``None`` if they can.""" |
| 200 | + if not get_accelerator().is_available() or get_accelerator().device_name() != "cuda": |
| 201 | + return "CUDA graphs require the CUDA accelerator" |
| 202 | + |
| 203 | + if zero_stage == 3: |
| 204 | + # Under ZeRO-3 the inference containers hold no persistent weights; the |
| 205 | + # parameters are gathered into fresh buffers for each generate call. A |
| 206 | + # graph would keep replaying whichever buffers existed at capture time, |
| 207 | + # which is silently wrong rather than merely slow. |
| 208 | + return "CUDA graphs are not supported with ZeRO stage 3" |
| 209 | + |
| 210 | + if config.release_inference_cache: |
| 211 | + # Releasing the workspace frees the buffers the graphs write into. |
| 212 | + return "CUDA graphs are not supported with release_inference_cache" |
| 213 | + |
| 214 | + if config.inference_tp_size > 1: |
| 215 | + return "CUDA graphs are not supported with inference_tp_size > 1" |
| 216 | + |
| 217 | + return None |
0 commit comments