ViperEkura
6054b8dbd4
- New CudaGraphContext class: warmup -> capture -> replay lifecycle - One graph per batch_size key, all inputs at fixed workspace addresses - Added position_ids buffer to InferenceWorkspace (required for graph capture) - Graph only activates when CUDA backend is the current backend - Default off (opt-in) due to slight numerical divergence in graph replay - Sampling stays outside the graph (torch.multinomial uses mutable RNG) - Resolved circular import: KVCache -> TYPE_CHECKING in attention_backend.py
102 lines
3.5 KiB
Python
102 lines
3.5 KiB
Python
"""CUDA-graph capture for the decode model-forward step.
|
|
|
|
Mirrors SGLang's cuda-graph manager: one graph per batch size. The graph
|
|
pair. The graph captures ``model.forward()`` with workspace-backed inputs
|
|
(all at fixed addresses). Before each replay the caller updates the input
|
|
buffer content in-place so the graph sees fresh data at the same tensor
|
|
addresses.
|
|
|
|
Only the model forward is captured — sampling runs outside the graph
|
|
(via ``torch.multinomial`` which consumes a mutable RNG state).
|
|
"""
|
|
|
|
import torch
|
|
from torch import Tensor
|
|
|
|
|
|
class CudaGraphContext:
|
|
"""CUDA-graph capture/replay for decode steps.
|
|
|
|
Parameters:
|
|
enabled: When ``False``, ``forward()`` always runs the live model
|
|
forward without capture/replay (graphs are cleared). Toggle at
|
|
runtime via the ``set_enabled()`` method.
|
|
|
|
Usage::
|
|
|
|
gctx = CudaGraphContext()
|
|
with torch.inference_mode():
|
|
outputs = gctx.forward(
|
|
model,
|
|
key=(batch_size,),
|
|
input_ids=workspace.input_ids[:b].unsqueeze(1),
|
|
input_mask=input_mask,
|
|
kv_cache=kv_cache,
|
|
position_ids=workspace.position_ids[:b].unsqueeze(1),
|
|
)
|
|
|
|
The first call at a given key runs *without* capture (warmup). The
|
|
second call captures the graph. Subsequent calls replay the captured
|
|
graph. A ``torch.cuda.synchronize()`` before capture drains in-flight
|
|
work so the graph trace is clean.
|
|
"""
|
|
|
|
def __init__(self, enabled: bool = False):
|
|
self._enabled = enabled
|
|
self._graphs: dict[tuple, torch.cuda.CUDAGraph] = {}
|
|
self._outputs: dict[tuple, dict[str, Tensor]] = {}
|
|
self._warmed: set[tuple] = set()
|
|
|
|
@property
|
|
def enabled(self) -> bool:
|
|
return self._enabled
|
|
|
|
def set_enabled(self, flag: bool):
|
|
"""Enable or disable CUDA-graph capture at runtime.
|
|
|
|
Disabling clears all captured graphs (frees GPU memory) and warmup
|
|
state. Re-enabling after disable starts fresh — graphs are
|
|
re-captured on the next warmup cycle.
|
|
"""
|
|
if flag == self._enabled:
|
|
return
|
|
self._enabled = flag
|
|
if not flag:
|
|
self._graphs.clear()
|
|
self._outputs.clear()
|
|
self._warmed.clear()
|
|
|
|
def forward(self, model, *, key, **kwargs) -> dict[str, Tensor]:
|
|
"""Run ``model(**kwargs)`` via graph replay or live forward.
|
|
|
|
Args:
|
|
model: callable, e.g. ``self.model.forward``.
|
|
key: ``(batch_size,)`` — the dispatch key (one graph per batch size).
|
|
**kwargs: arguments forwarded to ``model``. All tensor arguments
|
|
must reside at stable addresses (workspace buffers).
|
|
|
|
Returns:
|
|
The dict produced by ``model(**kwargs)``, e.g.
|
|
``{"logits": ..., "h0": ...}``.
|
|
"""
|
|
if not self._enabled:
|
|
self._outputs[key] = model(**kwargs)
|
|
return self._outputs[key]
|
|
|
|
if key in self._graphs:
|
|
self._graphs[key].replay()
|
|
elif key in self._warmed:
|
|
torch.cuda.synchronize()
|
|
graph = torch.cuda.CUDAGraph()
|
|
with torch.cuda.graph(graph):
|
|
self._outputs[key] = model(**kwargs)
|
|
self._graphs[key] = graph
|
|
self._warmed.discard(key)
|
|
else:
|
|
self._warmed.add(key)
|
|
self._outputs[key] = model(**kwargs)
|
|
return self._outputs[key]
|
|
|
|
def has_graph(self, key: tuple) -> bool:
|
|
return key in self._graphs
|