279 lines
11 KiB
Python
279 lines
11 KiB
Python
|
import os
|
||
|
|
|
||
|
|
import torch
|
||
|
|
|
||
|
|
ENABLE_CUDA_GRAPH = os.environ.get("ENABLE_CUDA_GRAPH", "True").lower() == "true"
|
||
|
|
# When using a single master buffer + multi-bucket graph, max batch must be pre-allocated;
|
||
|
|
# otherwise resize/reallocation after capture invalidates addresses recorded by the graph.
|
||
|
|
CUDA_GRAPH_MAX_BS = int(os.environ.get("CUDA_GRAPH_MAX_BS", "128"))
|
||
|
|
_SHAPE_GUARD_FASTPATH = (
|
||
|
|
os.environ.get("CUDAGRAPH_DISABLE_SHAPE_GUARD_FASTPATH", "0") != "1"
|
||
|
|
)
|
||
|
|
|
||
|
|
|
||
|
|
class CUDAGraph_Wrapper:
|
||
|
|
"""
|
||
|
|
Bucketed CUDA Graph wrapper:
|
||
|
|
- When batch size changes, select/reuse a CUDAGraph by bucket_size
|
||
|
|
- For batch_size < bucket_size: copy inputs into static buffer [:B], return outputs [:B] after replay
|
||
|
|
"""
|
||
|
|
|
||
|
|
def __init__(
|
||
|
|
self,
|
||
|
|
model,
|
||
|
|
warm_up_times: int = 3,
|
||
|
|
enable: bool | None = None,
|
||
|
|
batch_size_key: str | None = None,
|
||
|
|
):
|
||
|
|
self.model = model
|
||
|
|
self.warm_up_times = warm_up_times
|
||
|
|
self._enable = ENABLE_CUDA_GRAPH if enable is None else bool(enable)
|
||
|
|
self._batch_size_key = batch_size_key or "suffix_inputs_embeds"
|
||
|
|
|
||
|
|
# graph cache: bucket_bs -> CUDAGraph
|
||
|
|
self.graphs: dict[int, torch.cuda.CUDAGraph] = {}
|
||
|
|
# Per bucket, store views into the master buffer (graph capture depends on these views' address/shape/stride)
|
||
|
|
self.static_inputs_map: dict[int, dict[str, torch.Tensor]] = {}
|
||
|
|
self.static_output_tensor: dict[int, torch.Tensor] = {}
|
||
|
|
self.graph_pool = None
|
||
|
|
|
||
|
|
# Master buffers (shared across all buckets)
|
||
|
|
self.graph_vars: dict[str, torch.Tensor] = {}
|
||
|
|
self._max_bs: int | None = None
|
||
|
|
|
||
|
|
# Record non-batch dimension shape signature (globally consistent) to avoid silent errors
|
||
|
|
self._shape_signature: dict[str, tuple] = {}
|
||
|
|
# Batch dimension index per input (default 0; e.g. suffix_position_ids batch is at dim=1)
|
||
|
|
self._batch_dim: dict[str, int] = {}
|
||
|
|
self._shape_guard_verified: dict[int, bool] = {}
|
||
|
|
|
||
|
|
def _get_batch_size(self, **kwargs) -> int:
|
||
|
|
key = self._batch_size_key
|
||
|
|
if key in kwargs:
|
||
|
|
v = kwargs[key]
|
||
|
|
if isinstance(v, torch.Tensor) and v.dim() >= 1:
|
||
|
|
return int(v.shape[0])
|
||
|
|
raise ValueError(f"CUDAGraph_Wrapper.forward requires {key} argument")
|
||
|
|
|
||
|
|
def _select_bucket_bs(self, bs: int) -> int:
|
||
|
|
"""
|
||
|
|
Bucket selection:
|
||
|
|
- 1, 2, 4, 8 use pow2 buckets
|
||
|
|
- >8 align to 16 (16, 32, 48, ...)
|
||
|
|
"""
|
||
|
|
if bs <= 1:
|
||
|
|
return 1
|
||
|
|
if bs <= 2:
|
||
|
|
return 2
|
||
|
|
if bs <= 4:
|
||
|
|
return 4
|
||
|
|
if bs <= 8:
|
||
|
|
return 8
|
||
|
|
return int(((bs + 15) // 16) * 16)
|
||
|
|
|
||
|
|
def _batch_dim_for_key(self, key: str, tensor: torch.Tensor) -> int:
|
||
|
|
# Minimal special-case handling for known inputs only
|
||
|
|
# suffix_position_ids: shape [3, B, T], batch at dim=1
|
||
|
|
if key == "suffix_position_ids" and tensor.dim() >= 2:
|
||
|
|
return 1
|
||
|
|
return 0
|
||
|
|
|
||
|
|
def _view_for_bucket(self, tensor: torch.Tensor, bucket_bs: int, batch_dim: int):
|
||
|
|
if tensor.dim() == 0:
|
||
|
|
return tensor
|
||
|
|
if batch_dim == 0:
|
||
|
|
return tensor[:bucket_bs]
|
||
|
|
if batch_dim == 1:
|
||
|
|
return tensor[:, :bucket_bs]
|
||
|
|
raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper")
|
||
|
|
|
||
|
|
def _copy_and_pad(
|
||
|
|
self,
|
||
|
|
master: torch.Tensor,
|
||
|
|
tensor: torch.Tensor,
|
||
|
|
bs: int,
|
||
|
|
bucket_bs: int,
|
||
|
|
batch_dim: int,
|
||
|
|
):
|
||
|
|
if tensor.dim() == 0:
|
||
|
|
master.copy_(tensor)
|
||
|
|
return
|
||
|
|
if batch_dim == 0:
|
||
|
|
master[:bs].copy_(tensor)
|
||
|
|
if bs < bucket_bs:
|
||
|
|
master[bs:bucket_bs].zero_()
|
||
|
|
return
|
||
|
|
if batch_dim == 1:
|
||
|
|
master[:, :bs].copy_(tensor)
|
||
|
|
if bs < bucket_bs:
|
||
|
|
master[:, bs:bucket_bs].zero_()
|
||
|
|
return
|
||
|
|
raise ValueError(f"Unsupported batch_dim={batch_dim} for cudagraph wrapper")
|
||
|
|
|
||
|
|
def _allocate_master_tensor(
|
||
|
|
self, key: str, tensor: torch.Tensor, max_bs: int, bs: int
|
||
|
|
) -> torch.Tensor:
|
||
|
|
"""
|
||
|
|
Allocate a master buffer for one input (single large buffer):
|
||
|
|
- Extend only the batch dimension to max_bs; other dims unchanged
|
||
|
|
- Copy the current bs valid range first; zero the rest (avoid stale values on replay)
|
||
|
|
"""
|
||
|
|
if tensor.dim() == 0:
|
||
|
|
return tensor.clone()
|
||
|
|
batch_dim = self._batch_dim_for_key(key, tensor)
|
||
|
|
shape = list(tensor.shape)
|
||
|
|
if batch_dim >= len(shape):
|
||
|
|
raise ValueError(
|
||
|
|
f"Invalid batch_dim={batch_dim} for key={key}, tensor.shape={tuple(tensor.shape)}"
|
||
|
|
)
|
||
|
|
shape[batch_dim] = max_bs
|
||
|
|
buf = torch.zeros(tuple(shape), device=tensor.device, dtype=tensor.dtype)
|
||
|
|
# Initial copy: pad to current bucket (at least cover the valid bs range)
|
||
|
|
self._copy_and_pad(buf, tensor, bs=bs, bucket_bs=bs, batch_dim=batch_dim)
|
||
|
|
return buf
|
||
|
|
|
||
|
|
def _ensure_master_buffers(self, bs: int, **kwargs):
|
||
|
|
"""
|
||
|
|
Ensure master buffers are allocated once.
|
||
|
|
- max_bs from CUDA_GRAPH_MAX_BS; if unset, use the first bucket_bs as max_bs
|
||
|
|
and disallow exceeding it later (resize would invalidate captured graphs).
|
||
|
|
"""
|
||
|
|
if self._max_bs is not None:
|
||
|
|
return
|
||
|
|
|
||
|
|
bucket_bs = self._select_bucket_bs(bs)
|
||
|
|
max_bs = CUDA_GRAPH_MAX_BS if CUDA_GRAPH_MAX_BS > 0 else bucket_bs
|
||
|
|
self._max_bs = int(max_bs)
|
||
|
|
|
||
|
|
for k, v in kwargs.items():
|
||
|
|
if not isinstance(v, torch.Tensor):
|
||
|
|
raise TypeError(
|
||
|
|
f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}"
|
||
|
|
)
|
||
|
|
batch_dim = self._batch_dim_for_key(k, v) if v.dim() > 0 else 0
|
||
|
|
self._batch_dim[k] = batch_dim
|
||
|
|
# Record non-batch dimensions (globally consistent)
|
||
|
|
if v.dim() == 0:
|
||
|
|
self._shape_signature[k] = tuple()
|
||
|
|
else:
|
||
|
|
sig = list(v.shape)
|
||
|
|
sig.pop(batch_dim)
|
||
|
|
self._shape_signature[k] = tuple(sig)
|
||
|
|
self.graph_vars[k] = self._allocate_master_tensor(
|
||
|
|
k, v, max_bs=self._max_bs, bs=bs
|
||
|
|
)
|
||
|
|
|
||
|
|
def _initialize_bucket(self, bucket_bs: int, bs: int, **kwargs):
|
||
|
|
assert self._max_bs is not None
|
||
|
|
self._shape_guard_verified[bucket_bs] = False
|
||
|
|
if bucket_bs > self._max_bs:
|
||
|
|
raise ValueError(
|
||
|
|
f"CUDAGraph bucket_bs({bucket_bs}) exceeds master max_bs({self._max_bs}). "
|
||
|
|
f"Set CUDA_GRAPH_MAX_BS >= {bucket_bs}."
|
||
|
|
)
|
||
|
|
|
||
|
|
print(f"[CUDAGraph] Initializing bucket_bs={bucket_bs} (current bs={bs}) ...")
|
||
|
|
|
||
|
|
# Build view map for this bucket (fixed address/shape/stride)
|
||
|
|
static_map: dict[str, torch.Tensor] = {}
|
||
|
|
for k, v in kwargs.items():
|
||
|
|
# Shape guard: non-batch dims must match
|
||
|
|
expected = self._shape_signature.get(k)
|
||
|
|
if v.dim() == 0:
|
||
|
|
got = tuple()
|
||
|
|
else:
|
||
|
|
batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v))
|
||
|
|
got_list = list(v.shape)
|
||
|
|
got_list.pop(batch_dim)
|
||
|
|
got = tuple(got_list)
|
||
|
|
if expected is not None and got != expected:
|
||
|
|
raise ValueError(
|
||
|
|
f"CUDAGraph bucket init: non-batch dim changed: key={k}, expected={expected}, got={got}"
|
||
|
|
)
|
||
|
|
master = self.graph_vars[k]
|
||
|
|
batch_dim = self._batch_dim.get(k, 0)
|
||
|
|
static_map[k] = self._view_for_bucket(
|
||
|
|
master, bucket_bs=bucket_bs, batch_dim=batch_dim
|
||
|
|
)
|
||
|
|
|
||
|
|
# Warmup (stabilize kernels/caches)
|
||
|
|
out = None
|
||
|
|
for _ in range(self.warm_up_times):
|
||
|
|
out = self.forward_naive(**static_map)
|
||
|
|
assert isinstance(out, torch.Tensor)
|
||
|
|
|
||
|
|
# Output master buffer (allocate once)
|
||
|
|
if "_outputs" not in self.graph_vars:
|
||
|
|
out_shape = (self._max_bs,) + tuple(out.shape[1:])
|
||
|
|
self.graph_vars["_outputs"] = torch.empty(
|
||
|
|
out_shape, device=out.device, dtype=out.dtype
|
||
|
|
)
|
||
|
|
|
||
|
|
self.static_inputs_map[bucket_bs] = static_map
|
||
|
|
self.static_output_tensor[bucket_bs] = self.graph_vars["_outputs"][:bucket_bs]
|
||
|
|
|
||
|
|
# capture
|
||
|
|
graph = torch.cuda.CUDAGraph()
|
||
|
|
with torch.cuda.graph(graph, pool=self.graph_pool):
|
||
|
|
self.static_output_tensor[bucket_bs].copy_(self.forward_naive(**static_map))
|
||
|
|
if self.graph_pool is None:
|
||
|
|
self.graph_pool = graph.pool()
|
||
|
|
self.graphs[bucket_bs] = graph
|
||
|
|
torch.cuda.synchronize()
|
||
|
|
print(f"[CUDAGraph] Bucket {bucket_bs} captured.")
|
||
|
|
|
||
|
|
def forward_naive(self, **kwargs):
|
||
|
|
return self.model(**kwargs)
|
||
|
|
|
||
|
|
def forward(self, **kwargs):
|
||
|
|
if not self._enable:
|
||
|
|
return self.forward_naive(**kwargs)
|
||
|
|
bs = self._get_batch_size(**kwargs)
|
||
|
|
bucket_bs = self._select_bucket_bs(bs)
|
||
|
|
|
||
|
|
# Initialize master buffers (first call)
|
||
|
|
self._ensure_master_buffers(bs=bs, **kwargs)
|
||
|
|
assert self._max_bs is not None
|
||
|
|
if bucket_bs > self._max_bs:
|
||
|
|
# Do not exceed master max_bs at runtime (resize would invalidate captured graphs)
|
||
|
|
return self.forward_naive(**kwargs)
|
||
|
|
|
||
|
|
# lazy capture for this bucket
|
||
|
|
if bucket_bs not in self.graphs:
|
||
|
|
self._initialize_bucket(bucket_bs=bucket_bs, bs=bs, **kwargs)
|
||
|
|
|
||
|
|
# Shape guard: non-batch dim change is unsafe for cudagraph; fallback
|
||
|
|
# fast-path: skip guard if already verified for this bucket (stable-state optimization)
|
||
|
|
if not (
|
||
|
|
_SHAPE_GUARD_FASTPATH and self._shape_guard_verified.get(bucket_bs, False)
|
||
|
|
):
|
||
|
|
for k, v in kwargs.items():
|
||
|
|
if not isinstance(v, torch.Tensor):
|
||
|
|
raise TypeError(
|
||
|
|
f"CUDAGraph_Wrapper only supports Tensor kwargs, got {k}={type(v)}"
|
||
|
|
)
|
||
|
|
expected = self._shape_signature.get(k)
|
||
|
|
if v.dim() == 0:
|
||
|
|
got = tuple()
|
||
|
|
else:
|
||
|
|
batch_dim = self._batch_dim.get(k, self._batch_dim_for_key(k, v))
|
||
|
|
got_list = list(v.shape)
|
||
|
|
got_list.pop(batch_dim)
|
||
|
|
got = tuple(got_list)
|
||
|
|
if expected is not None and got != expected:
|
||
|
|
return self.forward_naive(**kwargs)
|
||
|
|
# Mark verified for fast-path on subsequent calls
|
||
|
|
self._shape_guard_verified[bucket_bs] = True
|
||
|
|
|
||
|
|
# Copy inputs into master buffer (per batch_dim, zero padding)
|
||
|
|
for k, tensor in kwargs.items():
|
||
|
|
master = self.graph_vars[k]
|
||
|
|
batch_dim = self._batch_dim.get(k, 0)
|
||
|
|
self._copy_and_pad(
|
||
|
|
master, tensor, bs=bs, bucket_bs=bucket_bs, batch_dim=batch_dim
|
||
|
|
)
|
||
|
|
|
||
|
|
graph = self.graphs[bucket_bs]
|
||
|
|
graph.replay()
|
||
|
|
return self.graph_vars["_outputs"][:bs]
|