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,33 +151,48 @@ 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.
-
-    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 _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:
-        return _current_backend.get()
-    except LookupError:
+                    _env_backend = AttentionBackendFactory.create(name)
+                except (ValueError, RuntimeError):
+                    _env_backend = None
+                _env_backend_name = name
+    return _env_backend
+
+
+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:
@@ -184,6 +201,22 @@ def get_backend() -> "AttentionBackend":
     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.
+    """
+    return (
+        _environment_backend()
+        or _current_backend.get()
+        or (_resolve_backend() if use_default else None)
+    )
+
+
 @contextmanager
 def attn_backend(backend: Union[str, ATTN_BACKEND, "AttentionBackend", type]):
     """Context manager to select an attention backend.
@@ -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,6 +77,23 @@ class InferenceScheduler:
             metrics=self._metrics,
         )
 
+        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,
@@ -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"):
-                step_out = self._executor.execute_decode(
-                    decoded, return_logprobs=return_logprobs
+        for backend, group in self._task_backend_groups(decoded):
+            backend_context = (
+                attn_backend(backend) if backend is not None else nullcontext()
             )
-            for t, out in zip(decoded, step_out):
+            with (
+                backend_context,
+                self._metrics.record([t.task_id for t in group], "decode"),
+            ):
+                step_out = self._executor.execute_decode(
+                    group, return_logprobs=return_logprobs
+                )
+            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,6 +234,7 @@ class InferenceScheduler:
     def _run_generation_loop(self):
         stop_ids = self._task_mgr.tokenizer.stop_ids
         try:
+            with self._backend_context():
                 while not self._stop_event.is_set():
                     finished = self._task_mgr.remove_finished_tasks(stop_ids)
                     for task in finished:
@@ -192,7 +253,9 @@ class InferenceScheduler:
                         candidates = self._task_mgr.pull_candidates(available)
                         failed = []
                         for task in candidates:
-                        if self._task_cache.task_alloc(task.task_id, task.prompt_ids):
+                            if self._task_cache.task_alloc(
+                                task.task_id, task.prompt_ids
+                            ):
                                 self._task_mgr.activate(task)
                             else:
                                 failed.append(task)
@@ -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,6 +385,7 @@ class InferenceScheduler:
         try:
             live = [t for t in tasks if t is not None]
 
+            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)]

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:

scripts/tools/benchmark.py

@@ -259,20 +259,19 @@ class GenerationBenchmark:
             max_seq_len=max_seq_len,
             cache=pool,
             enable_cuda_graph=self.cuda_graph,
+            backend=self.backend,
         )
         prompts = [prompt] * batch_size
 
         try:
             # Capture graphs and populate the allocator before timing. The
             # first request also includes model/scheduler startup effects.
-            with attn_backend(self.backend):
             engine.generate(prompts, max_tokens=gen_length, temperature=0.0)
             if self.device.startswith("cuda"):
                 torch.cuda.synchronize()
 
             t0 = time.perf_counter()
             for _ in range(num_trials):
-                with attn_backend(self.backend):
                 engine.generate(prompts, max_tokens=gen_length, temperature=0.0)
             if self.device.startswith("cuda"):
                 torch.cuda.synchronize()
@@ -291,7 +290,8 @@ class GenerationBenchmark:
                 "benchmark_type": "engine_decode",
                 "num_trials": num_trials,
                 "prompt_length": prompt_tokens,
-                "engine": True,
+                "backend": engine.backend_name,
+                "cuda_graph": engine.cuda_graph_enabled,
             },
         )
 

tests/extension/test_backend.py

@@ -42,6 +42,20 @@ def test_attn_backend_context_with_registered_name():
     assert get_backend() is default
 
 
+def test_backend_can_read_only_context_selection():
+    assert get_backend(use_default=False) is None
+    with attn_backend("cuda") as backend:
+        assert get_backend(use_default=False) is backend
+    assert get_backend(use_default=False) is None
+
+
+def test_environment_backend_overrides_context(monkeypatch):
+    monkeypatch.setenv("ASTR_BACKEND", "torch_native")
+    with attn_backend("cuda"):
+        assert type(get_backend()).__name__ == "TorchNativeBackend"
+        assert type(get_backend(use_default=False)).__name__ == "TorchNativeBackend"
+
+
 def test_attention_backend_factory_lists_builtin_backends():
     assert AttentionBackendFactory.list_registered() == [
         "cuda",

tests/inference/test_engine.py

@@ -3,6 +3,7 @@
 import threading
 from unittest.mock import MagicMock, patch
 
+from astrai.extension import TorchNativeBackend, attn_backend
 from astrai.inference import STOP
 from astrai.inference.engine import GenerateResult, InferenceEngine
 
@@ -199,3 +200,40 @@ def test_engine_generate_zero_max_tokens_stream_is_empty():
         eng = InferenceEngine(mock_model, mock_tokenizer, max_batch_size=1)
         assert list(eng.generate("hello", stream=True, max_tokens=0)) == []
         instance.add_task.assert_not_called()
+
+
+def test_engine_passes_backend_to_scheduler():
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        InferenceEngine(
+            mock_model,
+            mock_tokenizer,
+            max_batch_size=1,
+            backend="torch_native",
+        )
+
+    assert MockSched.call_args.kwargs["backend"] == "torch_native"
+
+
+def test_generate_captures_calling_backend_context():
+    mock_model = MagicMock()
+    mock_tokenizer = MagicMock()
+    captured = []
+
+    with patch("astrai.inference.engine.InferenceScheduler") as MockSched:
+        instance = MockSched.return_value
+
+        def fake_add(prompt, **kwargs):
+            captured.append(kwargs["backend"])
+            kwargs["stream_callback"](STOP)
+            return "task"
+
+        instance.add_task.side_effect = fake_add
+        engine = InferenceEngine(mock_model, mock_tokenizer)
+        with attn_backend("torch_native"):
+            assert engine.generate("hello") == ""
+
+    assert len(captured) == 1
+    assert isinstance(captured[0], TorchNativeBackend)

tests/inference/test_scheduler.py

@@ -7,7 +7,9 @@ from unittest.mock import MagicMock, patch
 import pytest
 import torch
 
+from astrai.extension import CudaBackend, TorchNativeBackend, get_backend
 from astrai.inference import InferenceScheduler
+from astrai.inference.metrics import MetricsCollector
 from astrai.inference.runtime.executor import DecodeSteadyState, Executor
 from astrai.inference.task import Task
 from astrai.model.transformer import AutoRegressiveLM
@@ -76,6 +78,69 @@ def test_scheduler_concurrent_add_task(mock_model_and_tokenizer):
     assert len(results["task_ids"]) == 50
 
 
+def test_generation_loop_activates_backend_in_worker_thread():
+    scheduler = object.__new__(InferenceScheduler)
+    scheduler._backend = TorchNativeBackend()
+    scheduler._stop_event = threading.Event()
+    scheduler._task_cache = MagicMock()
+
+    observed = []
+    task_mgr = MagicMock()
+    task_mgr.tokenizer.stop_ids = [0]
+    task_mgr.remove_finished_tasks.return_value = []
+    task_mgr.get_active_tasks.return_value = []
+    task_mgr.max_batch_size = 1
+    task_mgr.pull_candidates.return_value = []
+    task_mgr.has_work.return_value = False
+
+    def observe_backend(*args, **kwargs):
+        observed.append(type(get_backend()))
+        scheduler._stop_event.set()
+
+    task_mgr.wait_for_tasks.side_effect = observe_backend
+    scheduler._task_mgr = task_mgr
+
+    thread = threading.Thread(target=scheduler._run_generation_loop)
+    thread.start()
+    thread.join(timeout=5)
+
+    assert not thread.is_alive()
+    assert observed == [TorchNativeBackend]
+
+
+def test_step_splits_decode_batch_by_request_backend():
+    scheduler = object.__new__(InferenceScheduler)
+    scheduler._task_cache = MagicMock()
+    scheduler._task_cache.task_extend.return_value = True
+    scheduler._metrics = MetricsCollector()
+    scheduler._executor = MagicMock()
+
+    observed = []
+
+    def execute(tasks, **kwargs):
+        observed.append((type(get_backend()), [task.task_id for task in tasks]))
+        return [1] * len(tasks)
+
+    scheduler._executor.execute_decode.side_effect = execute
+
+    torch_task = Task("torch", [1], backend=TorchNativeBackend())
+    cuda_task = Task("cuda", [1], backend=CudaBackend())
+    for task in (torch_task, cuda_task):
+        task.input_tokens = 1
+        task.output_ids = [1]
+        task.mark_prefill_done()
+        scheduler._metrics.register(task.task_id)
+
+    produced, aborted = scheduler._step([torch_task, cuda_task])
+
+    assert aborted == []
+    assert produced == [torch_task, cuda_task]
+    assert observed == [
+        (TorchNativeBackend, ["torch"]),
+        (CudaBackend, ["cuda"]),
+    ]
+
+
 def test_scheduler_concurrent_add_remove_task(mock_model_and_tokenizer):
     """Test concurrent add and remove task operations."""
     mock_model, mock_tokenizer = mock_model_and_tokenizer