Skip to content

CUDA Graphs

CUDA-graphed wrapper for stateful models.

GraphedStatefulModel 

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

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.

Parameters:

Name Type Description Default
model Module

stateful model returning (output, state)

 required
num_warmup_iters int

warmup iterations before graph capture

 11
Source code in tsfast/models/_core/cudagraph.py 
42
43
44
45
46
47
48
49
def __init__(self, model: nn.Module, num_warmup_iters: int = 11):
    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

reset_graph 

reset_graph()

Clear captured graph for re-capture on next forward call.

Source code in tsfast/models/_core/cudagraph.py 
51
52
53
54
55
56
def reset_graph(self):
    """Clear captured graph for re-capture on next forward call."""
    self._graphed = None
    self._spec = None
    self._zero_flat = None
    self._graphed_shape = None