ViperEkura/AstrAI
1
0
Fork 0
Code Issues Pull Requests Packages Projects Releases Wiki Activity

feat: add CUDA-graph capture for decode forward

Browse Source 
- 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
This commit is contained in:
ViperEkura
2026-08-06 19:57:12 +08:00
parent 6f67ba8942
commit 6054b8dbd4
4 changed files with 165 additions and 32 deletions
Download Patch File Download Diff File Expand all files Collapse all files
astrai/extension/attention_backend.py
+16 -14
View File
@@ -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,