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
This commit is contained in:
2026-08-09 13:32:40 +08:00
parent c1d05ae11d
commit 47b3ed4e44
9 changed files with 334 additions and 110 deletions
+14
View File
@@ -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",