feat: propagate attention backend across scheduler threads

- InferenceEngine/Scheduler accept an explicit backend
- capture request-level attn_backend context onto Task
- split prefill/decode batches by backend instance
- ASTR_BACKEND env overrides ContextVar as process-wide policy
- report resolved backend and CUDA-graph state in benchmark

astrai/extension/attention_backend.py

@@ -56,8 +56,10 @@ if TYPE_CHECKING:
 
 _default_backend: Optional["AttentionBackend"] = None
 _default_backend_lock = threading.Lock()
-_current_backend: contextvars.ContextVar["AttentionBackend"] = contextvars.ContextVar(
-    "attn_backend"
+_env_backend_name: Optional[str] = None
+_env_backend: Optional["AttentionBackend"] = None
+_current_backend: contextvars.ContextVar[Optional["AttentionBackend"]] = (
+    contextvars.ContextVar("attn_backend", default=None)
 )
 
 
@@ -149,39 +151,70 @@ def _backend_supports(
 def _resolve_default_backend() -> "AttentionBackend":
     """Pick the highest-priority available backend (cuda -> flash -> torch).
 
-    Set ``ASTR_BACKEND`` to override: ``ASTR_BACKEND=cuda``, ``torch_native``,
-    or ``flash``.  The value is the registered name (same as the
-    ``ATTN_BACKEND`` enum value).
-
     Resolved lazily on first ``get_backend()`` and cached.  Per-call
     capability fallback happens in ``attention()``, so the default is
     safe for training and fp32 models.
     """
-    forced = os.environ.get("ASTR_BACKEND", "").strip().lower()
-    if forced:
-        try:
-            return AttentionBackendFactory.create(forced)
-        except (ValueError, RuntimeError):
-            pass
     return _priority_backends()[0]
 
 
-def get_backend() -> "AttentionBackend":
-    """Return the active backend for the current thread/context.
+def _environment_backend() -> Optional["AttentionBackend"]:
+    """Resolve the process-wide ``ASTR_BACKEND`` override, if configured."""
+    global _env_backend, _env_backend_name
+    name = os.environ.get("ASTR_BACKEND", "").strip().lower()
+    if not name:
+        return None
+    if name != _env_backend_name:
+        with _default_backend_lock:
+            if name != _env_backend_name:
+                try:
+                    _env_backend = AttentionBackendFactory.create(name)
+                except (ValueError, RuntimeError):
+                    _env_backend = None
+                _env_backend_name = name
+    return _env_backend
 
-    Falls back to the highest-priority available backend (cuda -> flash ->
-    torch_native) when no backend has been activated via ``with``.  Set
-    ``ASTR_BACKEND`` to override the default.
+
+def _resolve_backend(
+    backend: Optional[Union[str, ATTN_BACKEND, "AttentionBackend", type]] = None,
+) -> "AttentionBackend":
+    """Resolve a backend configuration, defaulting to the process policy."""
+    if backend is not None:
+        if isinstance(backend, ATTN_BACKEND):
+            return AttentionBackendFactory.create(backend.value)
+        if isinstance(backend, str):
+            return AttentionBackendFactory.create(backend)
+        if isinstance(backend, type) and issubclass(backend, AttentionBackend):
+            return backend()
+        if isinstance(backend, AttentionBackend):
+            return backend
+        raise TypeError(
+            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
+            f"or instance, got {type(backend).__name__}"
+        )
+
+    global _default_backend
+    if _default_backend is None:
+        with _default_backend_lock:
+            if _default_backend is None:
+                _default_backend = _resolve_default_backend()
+    return _default_backend
+
+
+def get_backend(
+    use_default: bool = True,
+) -> Optional["AttentionBackend"]:
+    """Return the context override, optionally falling back to the process default.
+
+    ``ASTR_BACKEND`` is a process-wide override and takes precedence over the
+    context value. Pass ``use_default=False`` at request submission to retain
+    only an environment override or the caller's :func:`attn_backend` value.
     """
-    try:
-        return _current_backend.get()
-    except LookupError:
-        global _default_backend
-        if _default_backend is None:
-            with _default_backend_lock:
-                if _default_backend is None:
-                    _default_backend = _resolve_default_backend()
-        return _default_backend
+    return (
+        _environment_backend()
+        or _current_backend.get()
+        or (_resolve_backend() if use_default else None)
+    )
 
 
 @contextmanager
@@ -200,20 +233,7 @@ def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
         with attn_backend(TorchNativeBackend()):
             ...
     """
-    if isinstance(backend, ATTN_BACKEND):
-        instance = AttentionBackendFactory.create(backend.value)
-    elif isinstance(backend, str):
-        instance = AttentionBackendFactory.create(backend)
-    elif isinstance(backend, type) and issubclass(backend, AttentionBackend):
-        instance = backend()
-    elif isinstance(backend, AttentionBackend):
-        instance = backend
-    else:
-        raise TypeError(
-            f"expected a registered name, ATTN_BACKEND, AttentionBackend type, "
-            f"or instance, "
-            f"got {type(backend).__name__}"
-        )
+    instance = _resolve_backend(backend)
     token = _current_backend.set(instance)
     try:
         yield instance
@@ -284,10 +304,7 @@ def attention(
     """
     backend = get_backend()
     if not _backend_supports(backend, q, kv_cache, attn_mask, is_causal):
-        try:
-            explicit = _current_backend.get()
-        except LookupError:
-            explicit = None
+        explicit = get_backend(use_default=False)
         if explicit is not None:
             raise RuntimeError(
                 f"Explicitly-set backend {type(backend).__name__} cannot "

astrai/inference/engine.py

@@ -8,6 +8,7 @@ from typing import Any, AsyncGenerator, Dict, Generator, List, Optional, Tuple,
 import torch
 import torch.nn as nn
 
+from astrai.extension import ATTN_BACKEND, AttentionBackend, get_backend
 from astrai.inference.cache import PagePool
 from astrai.inference.scheduler import InferenceScheduler
 from astrai.inference.task import STOP
@@ -75,6 +76,7 @@ class InferenceEngine:
         max_seq_len: Optional[int] = None,
         cache: Optional[PagePool] = None,
         enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
     ):
         self.model = model
         self.tokenizer = tokenizer
@@ -85,6 +87,7 @@ class InferenceEngine:
             max_seq_len=max_seq_len,
             cache=cache,
             enable_cuda_graph=enable_cuda_graph,
+            backend=backend,
         )
 
         self.scheduler.start()
@@ -174,6 +177,7 @@ class InferenceEngine:
         rep_window: int,
     ) -> Union[Generator, str, List[str]]:
         n = len(prompts)
+        request_backend = get_backend(use_default=False)
         result = GenerateResult(count=n)
         task_ids = [
             self.scheduler.add_task(
@@ -184,6 +188,7 @@ class InferenceEngine:
                 top_k=top_k,
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
+                backend=request_backend,
                 stream_callback=lambda token, idx=i: result.append(token, idx),
             )
             for i, p in enumerate(prompts)
@@ -223,6 +228,14 @@ class InferenceEngine:
     def get_stats(self) -> Dict[str, Any]:
         return self.scheduler.get_stats()
 
+    @property
+    def backend_name(self) -> str:
+        return self.scheduler.backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self.scheduler.cuda_graph_enabled
+
     def shutdown(self):
         self.scheduler.stop()
         if torch.cuda.is_available():

astrai/inference/runtime/executor.py

@@ -8,9 +8,7 @@ import torch
 from torch import Tensor
 
 from astrai.extension.attention_backend import (
-    ATTN_BACKEND,
     CudaBackend,
-    attn_backend,
     get_backend,
 )
 from astrai.inference.cache import PagePool, TaskCacheManager
@@ -153,7 +151,6 @@ def _warmup_cuda_graphs(
 
         with (
             torch.inference_mode(),
-            attn_backend(ATTN_BACKEND.CUDA),
             timed(f"warmup decode b={b}", logger),
         ):
             for step in range(2):
@@ -208,7 +205,10 @@ class Executor:
         config = model.config
         max_q_heads = config.num_attention_heads
         head_dim = config.hidden_size // config.num_attention_heads
-        self._graph_supported = CudaBackend.supports(head_dim=head_dim)
+        backend = get_backend()
+        self._graph_supported = backend.supports_graph() and CudaBackend.supports(
+            head_dim=head_dim
+        )
         self._workspace = InferenceWorkspace(
             max_batch_size=kv_cache.max_batch_size,
             max_seq_len=kv_cache.max_seq_len,
@@ -240,6 +240,10 @@ class Executor:
             device=self.device,
         )
 
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._graph_ctx.enabled and self._graph_supported
+
     def _sample_logits(
         self,
         logits: Tensor,

astrai/inference/scheduler.py

@@ -1,10 +1,17 @@
 import logging
 import threading
 import uuid
-from typing import Any, Dict, List, Optional, Tuple
+from contextlib import nullcontext
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 
+from astrai.extension import (
+    ATTN_BACKEND,
+    AttentionBackend,
+    attn_backend,
+    get_backend,
+)
 from astrai.inference.cache import PagePool, TaskCacheManager
 from astrai.inference.metrics import MetricsCollector
 from astrai.inference.runtime.executor import Executor
@@ -28,6 +35,7 @@ class InferenceScheduler:
         dtype: Optional[torch.dtype] = None,
         cache: Optional[PagePool] = None,
         enable_cuda_graph: bool = True,
+        backend: Optional[Union[str, ATTN_BACKEND, AttentionBackend, type]] = None,
     ):
         config = model.config
 
@@ -69,14 +77,31 @@ class InferenceScheduler:
             metrics=self._metrics,
         )
 
-        self._executor = Executor(
-            model=model,
-            kv_cache=self._cache,
-            task_cache=self._task_cache,
-            device=self.device,
-            dtype=self.dtype,
-            enable_cuda_graph=enable_cuda_graph,
-        )
+        if backend is None:
+            self._backend = None
+            default_backend = get_backend()
+            self._backend_name = type(default_backend).__name__
+            with attn_backend(default_backend):
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
+        else:
+            with attn_backend(backend):
+                self._backend = get_backend()
+                self._backend_name = type(self._backend).__name__
+                self._executor = Executor(
+                    model=model,
+                    kv_cache=self._cache,
+                    task_cache=self._task_cache,
+                    device=self.device,
+                    dtype=self.dtype,
+                    enable_cuda_graph=enable_cuda_graph,
+                )
 
         self._stop_event = threading.Event()
         self._loop_thread: Optional[threading.Thread] = None
@@ -91,6 +116,26 @@ class InferenceScheduler:
     def get_stats(self) -> Dict[str, Any]:
         return self._task_mgr.get_stats()
 
+    @property
+    def backend_name(self) -> str:
+        return self._backend_name
+
+    @property
+    def cuda_graph_enabled(self) -> bool:
+        return self._executor.cuda_graph_enabled
+
+    def _backend_context(self):
+        if self._backend is None:
+            return nullcontext()
+        return attn_backend(self._backend)
+
+    @staticmethod
+    def _task_backend_groups(tasks: List[Task]):
+        groups = {}
+        for task in tasks:
+            groups.setdefault(task.backend, (task.backend, []))[1].append(task)
+        return groups.values()
+
     def _step(
         self, tasks: List[Task], return_logprobs: bool = False
     ) -> Tuple[List[Task], List[Task]]:
@@ -121,15 +166,24 @@ class InferenceScheduler:
             for t in to_prefill:
                 t.input_tokens = len(t.prompt_ids)
 
-            groups: Dict[Tuple[int, int], List[Task]] = {}
+            groups: Dict[Tuple[int, int, Optional[AttentionBackend]], List[Task]] = {}
             for t in to_prefill:
                 start_pos = min(
                     self._task_cache.task_cached(t.task_id), len(t.prompt_ids) - 1
                 )
-                groups.setdefault((len(t.prompt_ids), start_pos), []).append(t)
+                groups.setdefault((len(t.prompt_ids), start_pos, t.backend), []).append(
+                    t
+                )
 
-            for (prompt_len, start_pos), group in groups.items():
-                with self._metrics.record([t.task_id for t in group], "prefill"):
+            for (prompt_len, start_pos, _), group in groups.items():
+                backend = group[0].backend
+                backend_context = (
+                    attn_backend(backend) if backend is not None else nullcontext()
+                )
+                with (
+                    backend_context,
+                    self._metrics.record([t.task_id for t in group], "prefill"),
+                ):
                     prefilled, step_out = self._executor.execute_prefill(
                         group, prompt_len, start_pos, return_logprobs=return_logprobs
                     )
@@ -158,12 +212,18 @@ class InferenceScheduler:
                 t.status = TaskStatus.ABORTED
                 aborted.append(t)
 
-        if decoded:
-            with self._metrics.record([t.task_id for t in decoded], "decode"):
+        for backend, group in self._task_backend_groups(decoded):
+            backend_context = (
+                attn_backend(backend) if backend is not None else nullcontext()
+            )
+            with (
+                backend_context,
+                self._metrics.record([t.task_id for t in group], "decode"),
+            ):
                 step_out = self._executor.execute_decode(
-                    decoded, return_logprobs=return_logprobs
+                    group, return_logprobs=return_logprobs
                 )
-            for t, out in zip(decoded, step_out):
+            for t, out in zip(group, step_out):
                 t.output_ids.append(out[0] if return_logprobs else out)
                 t.output_tokens += 1
                 t.advance_kv()
@@ -174,49 +234,52 @@ class InferenceScheduler:
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
         try:
-            while not self._stop_event.is_set():
-                finished = self._task_mgr.remove_finished_tasks(stop_ids)
-                for task in finished:
-                    if task.status == TaskStatus.FINISHED:
-                        self._task_cache.task_record_hashes(
-                            task.task_id,
-                            self._task_cache.task_cacheable_ids(
-                                task.task_id, task.prompt_ids, task.output_ids
-                            ),
-                        )
-                    self._task_cache.task_free(task.task_id)
+            with self._backend_context():
+                while not self._stop_event.is_set():
+                    finished = self._task_mgr.remove_finished_tasks(stop_ids)
+                    for task in finished:
+                        if task.status == TaskStatus.FINISHED:
+                            self._task_cache.task_record_hashes(
+                                task.task_id,
+                                self._task_cache.task_cacheable_ids(
+                                    task.task_id, task.prompt_ids, task.output_ids
+                                ),
+                            )
+                        self._task_cache.task_free(task.task_id)
 
-                active = self._task_mgr.get_active_tasks()
-                available = self._task_mgr.max_batch_size - len(active)
-                if available > 0:
-                    candidates = self._task_mgr.pull_candidates(available)
-                    failed = []
-                    for task in candidates:
-                        if self._task_cache.task_alloc(task.task_id, task.prompt_ids):
-                            self._task_mgr.activate(task)
-                        else:
-                            failed.append(task)
-                    if failed:
-                        self._task_mgr.return_to_waiting(failed)
+                    active = self._task_mgr.get_active_tasks()
+                    available = self._task_mgr.max_batch_size - len(active)
+                    if available > 0:
+                        candidates = self._task_mgr.pull_candidates(available)
+                        failed = []
+                        for task in candidates:
+                            if self._task_cache.task_alloc(
+                                task.task_id, task.prompt_ids
+                            ):
+                                self._task_mgr.activate(task)
+                            else:
+                                failed.append(task)
+                        if failed:
+                            self._task_mgr.return_to_waiting(failed)
 
-                if not self._task_mgr.has_work():
-                    self._task_mgr.wait_for_tasks(timeout=1.0)
-                    continue
+                    if not self._task_mgr.has_work():
+                        self._task_mgr.wait_for_tasks(timeout=1.0)
+                        continue
 
-                active = self._task_mgr.get_active_tasks()
+                    active = self._task_mgr.get_active_tasks()
 
-                decoded, aborted = self._step(active)
+                    decoded, aborted = self._step(active)
 
-                for t in aborted:
-                    self._task_mgr.invoke_callback(t.task_id, STOP)
-
-                for t in decoded:
-                    new_text = t.decode_new_token(self._task_mgr.tokenizer)
-                    if new_text:
-                        self._task_mgr.invoke_callback(t.task_id, new_text)
-                    if t.is_finished(stop_ids):
+                    for t in aborted:
                         self._task_mgr.invoke_callback(t.task_id, STOP)
 
+                    for t in decoded:
+                        new_text = t.decode_new_token(self._task_mgr.tokenizer)
+                        if new_text:
+                            self._task_mgr.invoke_callback(t.task_id, new_text)
+                        if t.is_finished(stop_ids):
+                            self._task_mgr.invoke_callback(t.task_id, STOP)
+
         except Exception as e:
             self._stop_event.set()
             logger.error(f"Scheduler loop crashed: {e}", exc_info=True)
@@ -286,6 +349,7 @@ class InferenceScheduler:
         """
         stop_ids = self._task_mgr.tokenizer.stop_ids
         seq_cap = self.max_seq_len
+        request_backend = get_backend(use_default=False)
 
         tasks: List[Task] = []
         for ids in prompt_ids_list:
@@ -309,6 +373,7 @@ class InferenceScheduler:
                 top_k=top_k,
                 frequency_penalty=frequency_penalty,
                 rep_window=rep_window,
+                backend=request_backend,
             )
             if not self._task_cache.task_alloc(task.task_id, task.prompt_ids):
                 tasks.append(None)
@@ -320,9 +385,10 @@ class InferenceScheduler:
         try:
             live = [t for t in tasks if t is not None]
 
-            while live:
-                decoded, _ = self._step(live, return_logprobs=return_logprobs)
-                live = [t for t in decoded if not t.is_finished(stop_ids)]
+            with self._backend_context():
+                while live:
+                    decoded, _ = self._step(live, return_logprobs=return_logprobs)
+                    live = [t for t in decoded if not t.is_finished(stop_ids)]
         finally:
             for t in tasks:
                 if t is not None:

astrai/inference/task.py

@@ -3,13 +3,16 @@ import time
 import uuid
 from collections import deque
 from enum import Enum
-from typing import Any, Callable, Deque, Dict, List, Optional
+from typing import TYPE_CHECKING, Any, Callable, Deque, Dict, List, Optional
 
 from tokenizers.decoders import DecodeStream
 
 from astrai.inference.metrics import MetricsCollector
 from astrai.tokenize.tokenizer import AutoTokenizer
 
+if TYPE_CHECKING:
+    from astrai.extension import AttentionBackend
+
 STOP = object()
 
 
@@ -62,6 +65,7 @@ class Task:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
     ):
         self.task_id = task_id
         self.prompt_ids = prompt_ids
@@ -71,6 +75,7 @@ class Task:
         self.top_k = top_k
         self.frequency_penalty = frequency_penalty
         self.rep_window = rep_window
+        self.backend = backend
 
         self.status = TaskStatus.PENDING
         self.output_ids: List[int] = []
@@ -152,6 +157,7 @@ class TaskManager:
         top_k: int = 50,
         frequency_penalty: float = 0.0,
         rep_window: int = 64,
+        backend: Optional["AttentionBackend"] = None,
         stream_callback: Optional[Callable[[str], None]] = None,
     ) -> str:
         task_id = f"task_{int(time.time())}_{uuid.uuid4().hex[:8]}"
@@ -173,6 +179,7 @@ class TaskManager:
             top_k=top_k,
             frequency_penalty=frequency_penalty,
             rep_window=rep_window,
+            backend=backend,
         )
 
         with self._lock: