refactor: simplify inference engine and backend dispatch

- merge _generate_streaming/_generate_non_streaming into single _generate() with stream flag
- delete dead GenerationRequest class and generate_with_request method
- inline _next_token helper into generate_async
- replace flash-attn double-checked locking with functools.lru_cache
- extract _write_and_gather_kv helper shared by TorchNative/FlashAttn backends
- inline _kv_cache_is_contiguous into its sole call site in FlashAttnBackend
- change default backend priority from flash>cuda>torch to cuda>flash>torch
- add ASTR_BACKEND env var to override default backend at resolve time
- add supports_graph() static method to AttentionBackend ABC, override in CudaBackend
- replace isinstance(get_backend(), CudaBackend) with get_backend().supports_graph() in executor
- add torch.cuda.is_available() guard to CudaBackend.supports()

tests/extension/test_backend.py

@@ -19,12 +19,13 @@ def test_default_backend_is_torch_native():
     """Default is the highest-priority available backend (flash > cuda > torch)."""
     from astrai.extension.attention_backend import (
         CudaBackend,
+        FlashAttnBackend,
         TorchNativeBackend,
         _resolve_default_backend,
     )
 
     backend = get_backend()
-    assert isinstance(backend, (CudaBackend, TorchNativeBackend))
+    assert isinstance(backend, (CudaBackend, FlashAttnBackend, TorchNativeBackend))
     assert isinstance(backend, type(_resolve_default_backend()))
 
 

tests/inference/test_protocol.py

@@ -8,7 +8,6 @@ import pytest
 from astrai.inference.api.anthropic import AnthropicResponseBuilder
 from astrai.inference.api.openai import OpenAIResponseBuilder
 from astrai.inference.api.protocol import GenContext, StopChecker, StopInfo
-from astrai.inference.engine import GenerationRequest
 
 
 def _make_ctx(**kwargs):
@@ -255,32 +254,3 @@ class TestAnthropicResponseBuilder:
         resp = builder.format_response(ctx, "full text", stop)
         assert resp["content"][0]["text"] == "full text"
         assert resp["stop_reason"] == "end_turn"
-
-
-class TestGenerationRequestValidation:
-    def test_valid_params(self):
-        gr = GenerationRequest(
-            messages=[{"role": "user", "content": "hi"}],
-            top_k=50,
-            top_p=0.9,
-            temperature=0.7,
-        )
-        assert gr.top_k == 50
-
-    def test_invalid_top_p_raises(self):
-        with pytest.raises(ValueError, match="top_p"):
-            GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_p=1.5)
-
-    def test_invalid_top_k_raises(self):
-        with pytest.raises(ValueError, match="top_k"):
-            GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_k=-1)
-
-    def test_invalid_temperature_raises(self):
-        with pytest.raises(ValueError, match="temperature"):
-            GenerationRequest(
-                messages=[{"role": "user", "content": "hi"}], temperature=-0.1
-            )
-
-    def test_top_k_zero_valid(self):
-        gr = GenerationRequest(messages=[{"role": "user", "content": "hi"}], top_k=0)
-        assert gr.top_k == 0