Skip to content

CUDA Graphs

CUDA-graphed wrapper for stateful models.

GraphedStatefulModel 

GraphedStatefulModel(model: Module, num_warmup_iters: int = 3)

Bases: Module

Wraps a stateful model with CUDA-graphed forward, same interface.

The model must return (output, state) from forward(). The CUDA graph is captured lazily on the first forward call. When input shapes change (e.g. different batch size at test time), falls back to eager execution automatically. A capture that fails outright also falls back to eager — permanently, with a RuntimeWarning naming the error — after restoring the CUDA state a broken capture leaves behind (current stream and RNG; see _recover_from_failed_capture).

Parameters:

Name Type Description Default
model Module

stateful model returning (output, state)

 required
num_warmup_iters int

warmup iterations before graph capture

 3
Source code in tsfast/models/_core/cudagraph.py 
69
70
71
72
73
74
75
76
77
def __init__(self, model: nn.Module, num_warmup_iters: int = 3):
    super().__init__()
    self.model = model
    self.num_warmup_iters = num_warmup_iters
    self._graphed = None
    self._spec: StateSpec | None = None
    self._zero_flat: Tensor | None = None
    self._graphed_shape: tuple[int, ...] | None = None
    self._capture_failed = False

reset_graph 

reset_graph()

Clear captured graph (and any capture failure) for re-capture on next forward call.

Source code in tsfast/models/_core/cudagraph.py 
79
80
81
82
83
84
85
def reset_graph(self):
    """Clear captured graph (and any capture failure) for re-capture on next forward call."""
    self._graphed = None
    self._spec = None
    self._zero_flat = None
    self._graphed_shape = None
    self._capture_failed = False