diff --git a/astrai/extension/attention_backend.py b/astrai/extension/attention_backend.py index 4448a0c..3dfe99a 100644 --- a/astrai/extension/attention_backend.py +++ b/astrai/extension/attention_backend.py @@ -34,7 +34,7 @@ import importlib import threading from abc import ABC, abstractmethod from contextlib import contextmanager -from typing import Optional, Union +from typing import TYPE_CHECKING, Optional, Union import torch import torch.nn.functional as F @@ -45,7 +45,9 @@ from astrai.extension.attention_ops import ( attn_paged_prefill, ) from astrai.factory import BaseFactory -from astrai.inference.core.cache import KVCache + +if TYPE_CHECKING: + from astrai.inference.core.cache import KVCache _current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar( "attn_backend" @@ -199,7 +201,7 @@ def attention( q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache] = None, + kv_cache: Optional["KVCache"] = None, layer_id: int = 0, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -255,7 +257,7 @@ class AttentionBackend(ABC): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -284,7 +286,7 @@ class AttentionBackend(ABC): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -297,7 +299,7 @@ class AttentionBackend(ABC): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -326,7 +328,7 @@ class TorchNativeBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -338,7 +340,7 @@ class TorchNativeBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -350,7 +352,7 @@ class TorchNativeBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -416,7 +418,7 @@ class CudaBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -451,7 +453,7 @@ class CudaBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -503,7 +505,7 @@ class FlashAttnBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -515,7 +517,7 @@ class FlashAttnBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, @@ -527,7 +529,7 @@ class FlashAttnBackend(AttentionBackend): q: Tensor, k: Tensor, v: Tensor, - kv_cache: Optional[KVCache], + kv_cache: Optional["KVCache"], layer_id: int, attn_mask: Optional[Tensor] = None, is_causal: bool = False, diff --git a/astrai/inference/core/executor.py b/astrai/inference/core/executor.py index 5f3bd3c..e19b6b0 100644 --- a/astrai/inference/core/executor.py +++ b/astrai/inference/core/executor.py @@ -5,7 +5,9 @@ from typing import List, Optional import torch from torch import Tensor +from astrai.extension.attention_backend import CudaBackend, get_backend from astrai.inference.core.cache import PagePool +from astrai.inference.core.graph import CudaGraphContext from astrai.inference.core.task import Task from astrai.inference.core.workspace import InferenceWorkspace from astrai.inference.sample import sample @@ -90,6 +92,12 @@ class Executor: dtype=self.dtype, ) + # CUDA-graph capture: one graph per (batch_size, total_len) key. + # The graph captures model.forward() with fixed-address workspace + # inputs. Before each replay, input content is updated in-place so + # the graph sees fresh token IDs / positions / KV metadata. + self._graph_ctx = CudaGraphContext() + def _sample_logits( self, logits: Tensor, @@ -204,42 +212,59 @@ class Executor: if not tasks: return [] - input_ids = self._workspace.fill_input_ids( + b = len(tasks) + ws = self._workspace + + # ---- pre-replay: update input buffers in-place ---- + + input_ids = ws.fill_input_ids( [t.output_ids[-1] if t.output_ids else t.prompt_ids[-1] for t in tasks] - ).unsqueeze(1) + ) task_ids = [t.task_id for t in tasks] + cur_positions = [t.next_pos for t in tasks] sig = tuple(task_ids) - cur_positions = [t.next_pos for t in tasks] cached = self._decode_cache if ( cached is not None and cached[0] == sig and cur_positions == [p + 1 for p in cached[1]] ): - _, _, info, position_ids = cached - position_ids += 1 - self._decode_cache = (sig, cur_positions, info, position_ids) + info = cached[2] + ws.position_ids[:b] += 1 + self._decode_cache = (sig, cur_positions, info) else: info = _build_sampling_batch_info(tasks, self.device) - position_ids = torch.tensor( - cur_positions, dtype=torch.long, device=self.device + ws.position_ids[:b].copy_( + torch.tensor(cur_positions, dtype=torch.long, device=self.device) ) - self._decode_cache = (sig, cur_positions, info, position_ids) + self._decode_cache = (sig, cur_positions, info) - total_len = max(t.next_pos for t in tasks) + 1 - input_mask = self._workspace.decode_mask(position_ids, total_len) + total_len = max(cur_positions) + 1 + input_mask = ws.decode_mask(ws.position_ids[:b], total_len) + + kv_cache = self.kv_cache.bind_tasks(task_ids, ws) + + # ---- forward (graph replay or live run + capture) ---- + + use_graph = ( + self._graph_ctx.enabled + and "cuda" in str(self.device) + and isinstance(get_backend(), CudaBackend) + ) + key = (b,) + if use_graph: + input_mask = ws.decode_mask(ws.position_ids[:b], ws.max_seq_len) with torch.inference_mode(): - outputs = self.model( - input_ids, + outputs = self._graph_ctx.forward( + self.model, + key=key, + input_ids=input_ids.unsqueeze(1), input_mask=input_mask, - kv_cache=self.kv_cache.bind_tasks( - task_ids, - self._workspace, - ), - position_ids=position_ids.unsqueeze(1), + kv_cache=kv_cache, + position_ids=ws.position_ids[:b].unsqueeze(1), ) logits = outputs["logits"][:, -1, :] diff --git a/astrai/inference/core/graph.py b/astrai/inference/core/graph.py new file mode 100644 index 0000000..8692ca1 --- /dev/null +++ b/astrai/inference/core/graph.py @@ -0,0 +1,101 @@ +"""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 diff --git a/astrai/inference/core/workspace.py b/astrai/inference/core/workspace.py index 4eb93a6..4bca49f 100644 --- a/astrai/inference/core/workspace.py +++ b/astrai/inference/core/workspace.py @@ -88,6 +88,11 @@ class InferenceWorkspace: (max_batch_size, 1), dtype=torch.long, device=device ) + # Per-step position IDs (must be at a fixed address for CUDA-graph capture). + self.position_ids = torch.empty( + (max_batch_size,), dtype=torch.long, device=device + ) + # Split-KV partial-result buffers for decode (persistent, one global # alloc per process — mirrors FlashInfer's workspace pattern). # Shape: [max_batch_size, max_q_heads, _MAX_SPLITS, head_dim] (o_part)