Dispatches generation forwards to a per-position CUDA graph. Args: forward: The original ``forward``, used both for capture and as the eager fallback. max_positions: Upper bound on how many decode positions may be captured, which bounds the memory the cac
| 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 |
no outgoing calls