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()
This commit is contained in:
2026-08-07 22:28:48 +08:00
parent 05739629fc
commit 02469887f5
7 changed files with 119 additions and 299 deletions
+2 -3
View File
@@ -5,7 +5,7 @@ Layers:
- api/: HTTP orchestration (ProtocolHandler, server)
- protocols/: Response builders (OpenAI, Anthropic)
- transport/: SSE transport utilities
- engine.py: Facade (InferenceEngine), Value Object (GenerationRequest)
- engine.py: Facade (InferenceEngine)
- sample.py: Strategy pattern (TemperatureStrategy, TopKStrategy, TopPStrategy, FrequencyPenaltyStrategy)
"""
@@ -42,7 +42,7 @@ from astrai.inference.core import (
TaskStatus,
page_hash,
)
from astrai.inference.engine import GenerationRequest, InferenceEngine
from astrai.inference.engine import InferenceEngine
from astrai.inference.sample import (
BaseSamplingStrategy,
FrequencyPenaltyStrategy,
@@ -55,7 +55,6 @@ from astrai.inference.sample import (
__all__ = [
"InferenceEngine",
"GenerationRequest",
"InferenceScheduler",
"Executor",
"STOP",